Search results

How to integrate Torch model with Flax nnx?

Status: Open.

#94 In google/torchax;

· jorgenin opened on May 11, 2026

How torchax handle tied_embedding?

Status: Open.

#68 In google/torchax;

· carlesoctav opened on Jan 22, 2026

train_llama_torchtitan example: Multiple issues running on TPU v4-32

Status: Open.

#67 In google/torchax;

· KarthikRevanuru opened on Jan 21, 2026

Convolution backward op missing

Status: Open.

#65 In google/torchax;

· samos123 opened on Jan 14, 2026

Unable to export quantized models to stablehlo

Status: Open.

#62 In google/torchax;

· Ha-Eyz opened on Dec 18, 2025

Torchax vs Torch2Jax ?

Status: Open.

#55 In google/torchax;

· adam-hartshorne opened on Nov 25, 2025

Dependency issue of 0.0.9

Status: Open.

#30 In google/torchax;

· wdhongtw opened on Nov 7, 2025

Buffers not moving to jax device

Status: Open.

#23 In google/torchax;

· mitkotak opened on Oct 29, 2025

Would nn.BatchNorm2d work as expected?

Status: Open.

#22 In google/torchax;

· lime-j opened on Oct 26, 2025

How do RNGs work in torchax

Status: Open.

#17 In google/torchax;

· carlesoctav opened on Oct 24, 2025

GPU out of memory during large model transfer to Jax device

Status: Open.

#16 In google/torchax;

· mozuysal-hubx opened on Oct 23, 2025

OperatorNotFound: Operator with name aten::_nested_tensor_from_mask_left_aligned has no lowering

Status: Open.

#12 In google/torchax;

· BioGeek opened on Oct 17, 2025