instadeepai/Mava

[FEATURE] Create an `add_batch_dim` method

Open

#1153 aperta il 9 gen 2025

Vedi su GitHub
 (0 commenti) (0 reazioni) (0 assegnatari)Python (122 fork)auto 404
enhancementgood first issue

Metriche repository

Star
 (922 star)
Metriche merge PR
 (Metriche PR in attesa)

Descrizione

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

Guida contributor

Direzione di ricerca
Crea una funzione di utilità in `utils/jax utils.py` che prende un pytree e restituisce lo stesso pytree con una nuova dimensione batch aggiunta usando `jax.tree.map(lambda x: x[jnp.newaxis], pytree)`. Assicurati di importare `jax.numpy as jnp` e `jax.tree util`.
Tech stack
python
Dominio
backend
Tipo issue
Funzionalità
DifficoltàQuanto dovrebbe essere impegnativa per un nuovo contributor.
1
Tempo stimatoTempo stimato per completare e verificare uno scope piccolo.
Meno di 1 ora
Stato attivitàQuanto è probabile che la issue sia ancora attiva e reviewable.
Nuova
ChiarezzaQuanto chiaramente la issue descrive problema, scope e risultato atteso.
Chiara
Prerequisiti
PythonJAX
Adatta ai principiantiQuanto la issue è adatta a contributor alla prima esperienza.
90