instadeepai/Mava

[FEATURE] Create an `add_batch_dim` method

Open

#1 153 ouverte le 9 janv. 2025

Voir sur GitHub
 (0 commentaires) (0 réactions) (0 assignés)Python (122 forks)auto 404
enhancementgood first issue

Métriques du dépôt

Stars
 (922 stars)
Métriques de merge PR
 (Métriques PR en attente)

Description

All over Mava we have statements like tree.map(lambda x: x[jnp.newaxis], pytree) I think it would make it a lot clearer if we had an add_batch_dim method in utils.jax_utils

Guide contributeur

Direction de recherche
Créez une fonction utilitaire dans `utils/jax utils.py` qui prend un pytree et retourne le même pytree avec une nouvelle dimension de lot ajoutée en utilisant `jax.tree.map(lambda x: x[jnp.newaxis], pytree)`. Assurez vous d'importer `jax.numpy as jnp` et `jax.tree util`.
Stack technique
python
Domaine
backend
Type d'issue
Fonctionnalité
DifficultéDifficulté estimée pour un nouveau contributeur, de 1 pour un très petit changement à 5 pour un travail expert.
1
Temps estiméFourchette de temps approximative pour investiguer, implémenter, tester et préparer une pull request.
Moins d'une heure
Statut d'activitéDisponibilité apparente de l'issue : fraîche, active, ancienne, bloquée ou en attente d'un mainteneur.
Récente
ClartéClarté avec laquelle l'issue explique le changement attendu, les critères d'acceptation et la prochaine étape.
Claire
Prérequis
PythonJAX
Accessibilité débutantScore de 1 à 100 estimant l'accessibilité de cette issue pour un premier contributeur.
90