Merge pull request #9052 from jpuigcerver:main

PiperOrigin-RevId: 430680329
jax-ml · Feb 24, 2022 · 3948fde · 3948fde
2 parents a9a827e + 86e8928
commit 3948fde
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/jax/_src/nn/initializers.py b/jax/_src/nn/initializers.py
@@ -32,6 +32,12 @@
 def zeros(key, shape, dtype=jnp.float_): return jnp.zeros(shape, dtypes.canonicalize_dtype(dtype))
 def ones(key, shape, dtype=jnp.float_): return jnp.ones(shape, dtypes.canonicalize_dtype(dtype))
 
+def constant(value, dtype=jnp.float_):
+  def init(key, shape, dtype=dtype):
+    dtype = dtypes.canonicalize_dtype(dtype)
+    return jnp.full(shape, value, dtype=dtype)
+  return init
+
 def uniform(scale=1e-2, dtype=jnp.float_):
   def init(key, shape, dtype=dtype):
     dtype = dtypes.canonicalize_dtype(dtype)

diff --git a/jax/nn/initializers.py b/jax/nn/initializers.py
@@ -19,6 +19,7 @@
 
 # flake8: noqa: F401
 from jax._src.nn.initializers import (
+  constant as constant,
   delta_orthogonal as delta_orthogonal,
   glorot_normal as glorot_normal,
   glorot_uniform as glorot_uniform,