jax-ml/jax · error · ValueError

orthogonal initializer requires at least a 2D shape

Error message

orthogonal initializer requires at least a 2D shape

What it means

jax.nn.initializers.orthogonal produces a (semi-)orthogonal matrix via QR/SVD, which requires at least 2 dimensions. Passing a 0D or 1D shape raises this ValueError.

Source

Thrown at jax/_src/nn/initializers.py:632

    An orthogonal initializer.

  Examples:

  >>> import jax, jax.numpy as jnp
  >>> initializer = jax.nn.initializers.orthogonal()
  >>> initializer(jax.random.key(42), (2, 3), jnp.float32)  # doctest: +SKIP
  Array([[ 3.9026976e-01,  7.2495741e-01, -5.6756169e-01],
         [ 8.8047469e-01, -4.7409311e-01, -1.3157725e-04]],            dtype=float32)
  """
  def init(key: Array,
           shape: core.Shape,
           dtype: DTypeLikeInexact | None = dtype,
           out_sharding: OutShardingType = None) -> Array:
    if out_sharding is not None:
      raise NotImplementedError
    dtype = dtypes.default_float_dtype() if dtype is None else dtype
    if len(shape) < 2:
      raise ValueError("orthogonal initializer requires at least a 2D shape")

    if any(dim == 0 for dim in shape):
      # empty shape
      return jnp.zeros(shape, dtype=dtype, out_sharding=out_sharding)

    n_rows, n_cols = math.prod(shape) // shape[column_axis], shape[column_axis]
    Q = random.orthogonal(key, n_rows, (), dtype, n_cols)
    Q = jnp.reshape(Q, tuple(np.delete(shape, column_axis)) + (shape[column_axis],))
    Q = jnp.moveaxis(Q, -1, column_axis)
    return jnp.array(scale, dtype) * Q
  return init

@export
def delta_orthogonal(
  scale: RealNumeric = 1.0,
  column_axis: int = -1,
  dtype: DTypeLikeInexact | None = None) -> Initializer:
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use delta_initializer() or zeros for 1D parameters
2. Reshape 1D params to (n, 1) only if mathematically appropriate (a single orthogonal column is just a unit vector)

Example fix

// before
b = jax.nn.initializers.orthogonal()(key, (128,))

// after
b = jax.nn.initializers.zeros(key, (128,))

def pick_initializer(shape):
    if len(shape) < 2:
        return jax.nn.initializers.zeros  # or normal
    return jax.nn.initializers.orthogonal()

def supports_orthogonal(shape) -> bool: return len(shape) >= 2

When it happens

Trigger: Calling orthogonal()(key, ()) or orthogonal()(key, (128,)); transposing a Linear layer's shape incorrectly so only one dim reaches the initializer.

Common situations: Initializing biases or embeddings-of-rank-1 with orthogonal; generic layer code reusing one initializer for both weights (2D) and biases (1D).

Related errors

• Can't compute input and output sizes of a {len(shape)}-dimen
• {name} ndim should be {len(shape)}, but got {t.ndim}
• {name} shape should be {shape}: but got {t.shape}
• The number of query heads must be a multiple of key/value he
• scaled_matmul requires all inputs to be 3-dimensional arrays

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/08d44819ecbf054a. Report an issue: GitHub.