try to fix memory error

divyashreepathihalli · divyashreepathihalli · commit f0b10ef6155b · 2025-06-18T21:27:13.000Z
diff --git a/keras/src/backend/jax/core.py b/keras/src/backend/jax/core.py
@@ -83,7 +83,7 @@ def __init__(
 
             # Initialize nnx.Variable first.
             # Determine the dtype for the placeholder.
-            _placeholder_value = jnp.zeros(
+            _placeholder_value = jax.ShapeDtypeStruct(
                 shape or (), dtype=standardize_dtype(dtype)
             )
 

Original file line number	Diff line number	Diff line change
`@@ -83,7 +83,7 @@ def __init__(`
`83`	`83`
`84`	`84`	`# Initialize nnx.Variable first.`
`85`	`85`	`# Determine the dtype for the placeholder.`
`86`		`- _placeholder_value = jnp.zeros(`
	`86`	`+ _placeholder_value = jax.ShapeDtypeStruct(`
`87`	`87`	`shape or (), dtype=standardize_dtype(dtype)`
`88`	`88`	`)`
`89`	`89`