fixed typo in test utils

denera · denera · commit 663a57429f12 · 2025-07-09T19:46:49.000Z
Signed-off-by: Alp Dener &lt;adener@nvidia.com&gt;
diff --git a/tests/jax/utils.py b/tests/jax/utils.py
@@ -33,7 +33,7 @@
 PrecisionLike = Union[
     None, str, lax.Precision, Tuple[str, str], Tuple[lax.Precision, lax.Precision]
 ]
-Initializer = Callable[[PRNGKey, Shape, DType], Array]
+Initializer = Callable[[PRNGKey, Shape, jnp.dtype], Array]
 
 # Enables verbose printing of tensor numerics for debug.
 NVTE_DEBUG_NUMERICS = bool(int(os.getenv("NVTE_DEBUG_NUMERICS", 0)))

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@`
`33`	`33`	`PrecisionLike = Union[`
`34`	`34`	`None, str, lax.Precision, Tuple[str, str], Tuple[lax.Precision, lax.Precision]`
`35`	`35`	`]`
`36`		`-Initializer = Callable[[PRNGKey, Shape, DType], Array]`
	`36`	`+Initializer = Callable[[PRNGKey, Shape, jnp.dtype], Array]`
`37`	`37`
`38`	`38`	`# Enables verbose printing of tensor numerics for debug.`
`39`	`39`	`NVTE_DEBUG_NUMERICS = bool(int(os.getenv("NVTE_DEBUG_NUMERICS", 0)))`