Rename SDS(vma: frozenset = ...) to SDS(manual_type: jax.sharding.ManualAxisType = ...)

yashk2810 · Flax Authors · commit a626c66eca9f · 2026-03-17T19:16:47.000-07:00
PiperOrigin-RevId: 885309371
diff --git a/flax/core/axes_scan.py b/flax/core/axes_scan.py
@@ -52,7 +52,8 @@ def build_shaped_array(x, batch_dim: bool = False) -> core.ShapedArray:
       shape=shape,
       dtype=jnp.result_type(x),
       sharding=sharding,
-      **{k: getattr(x, k) for k in ["weak_type", "vma"] if hasattr(x, k)},
+      **{k: getattr(x, k) for k in ["weak_type", "manual_type"]
+         if hasattr(x, k)},
   )
 
 

Original file line number	Diff line number	Diff line change
`@@ -52,7 +52,8 @@ def build_shaped_array(x, batch_dim: bool = False) -> core.ShapedArray:`
`52`	`52`	`shape=shape,`
`53`	`53`	`dtype=jnp.result_type(x),`
`54`	`54`	`sharding=sharding,`
`55`		`- **{k: getattr(x, k) for k in ["weak_type", "vma"] if hasattr(x, k)},`
	`55`	`+ **{k: getattr(x, k) for k in ["weak_type", "manual_type"]`
	`56`	`+ if hasattr(x, k)},`
`56`	`57`	`)`
`57`	`58`
`58`	`59`