Repository: google/jax

Issue:
Rename jax.nn.normalize to jax.nn.standardize
As discussed: the jax.nn.normalize function actually standardizes by default, so we should name it accordingly.

Nobody reading `jax.nn.normalize(x)` expects the result not to have norm-1, and [PyTorch agrees with this too](https://pytorch.org/docs/stable/generated/torch.nn.functional.normalize.html).

The proposal here is just to rename to avoid user bugs. This would also open up the name `jax.nn.normalize` for a PyTorch-like implementation, but that should be as a follow-up, if desired.


Discussion:
Hrm, I guess we'd need to do a deprecation cycle for this. The steps may be:
1. in the first PR, rename `jax.nn.normalize` -> `jax.nn.standardize`, but also define `jax.nn.normalize` as an alias;
2. then, add a DeprecationWarning to `jax.nn.normalize` which points to the new API, and in the same CL update all Google-internal users to call `jax.nn.standardize`;
3. make a jax release;
4. wait [three months](https://jax.readthedocs.io/en/latest/api_compatibility.html);
5. either make `jax.nn.normalize` raise an error (instead of a warning) and wait some time, or just delete it;
6. optionally, repurpose `jax.nn.normalize`.

The first two steps could be done in rapid succession.

Anyone want to take it on? 😄 
I had a go at implementing this in https://github.com/google/jax/pull/9968. Let me know if I got it wrong or any suggested changes.

Resolving patch:
diff --git a/jax/_src/nn/functions.py b/jax/_src/nn/functions.py
--- a/jax/_src/nn/functions.py
+++ b/jax/_src/nn/functions.py
@@ -16,6 +16,7 @@
 
 from functools import partial
 import operator
+import warnings
 import numpy as np
 from typing import Any, Optional, Tuple, Union
 
@@ -333,7 +334,7 @@ def softmax(x: Array,
   return unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
 
 @partial(jax.jit, static_argnames=("axis",))
-def normalize(x: Array,
+def standardize(x: Array,
               axis: Optional[Union[int, Tuple[int, ...]]] = -1,
               mean: Optional[Array] = None,
               variance: Optional[Array] = None,
@@ -351,6 +352,15 @@ def normalize(x: Array,
         jnp.square(x), axis, keepdims=True, where=where) - jnp.square(mean)
   return (x - mean) * lax.rsqrt(variance + epsilon)
 
+def normalize(x: Array,
+              axis: Optional[Union[int, Tuple[int, ...]]] = -1,
+              mean: Optional[Array] = None,
+              variance: Optional[Array] = None,
+              epsilon: Array = 1e-5,
+              where: Optional[Array] = None) -> Array:
+  r"""Normalizes an array by subtracting ``mean`` and dividing by :math:`\sqrt{\mathrm{variance}}`."""
+  warnings.warn("jax.nn.normalize will be deprecated. Use jax.nn.standardize instead.", DeprecationWarning)
+  return standardize(x, axis, mean, variance, epsilon, where)
 
 @partial(jax.jit, static_argnames=("num_classes", "dtype", "axis"))
 def _one_hot(x: Array, num_classes: int, *,
diff --git a/jax/example_libraries/stax.py b/jax/example_libraries/stax.py
--- a/jax/example_libraries/stax.py
+++ b/jax/example_libraries/stax.py
@@ -34,7 +34,7 @@
 import jax.numpy as jnp
 
 from jax.nn import (relu, log_softmax, softmax, softplus, sigmoid, elu,
-                    leaky_relu, selu, gelu, normalize)
+                    leaky_relu, selu, gelu, standardize)
 from jax.nn.initializers import glorot_normal, normal, ones, zeros
 
 # aliases for backwards compatibility
@@ -137,7 +137,7 @@ def apply_fun(params, x, **kwargs):
     # TODO(phawkins): jnp.expand_dims should accept an axis tuple.
     # (https://github.com/numpy/numpy/issues/12290)
     ed = tuple(None if i in axis else slice(None) for i in range(jnp.ndim(x)))
-    z = normalize(x, axis, epsilon=epsilon)
+    z = standardize(x, axis, epsilon=epsilon)
     if center and scale: return gamma[ed] * z + beta[ed]
     if center: return z + beta[ed]
     if scale: return gamma[ed] * z
diff --git a/jax/nn/__init__.py b/jax/nn/__init__.py
--- a/jax/nn/__init__.py
+++ b/jax/nn/__init__.py
@@ -32,6 +32,7 @@
   log_softmax as log_softmax,
   logsumexp as logsumexp,
   normalize as normalize,
+  standardize as standardize,
   one_hot as one_hot,
   relu as relu,
   relu6 as relu6,


Test patch:
diff --git a/tests/nn_test.py b/tests/nn_test.py
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -126,13 +126,13 @@ def testSoftmaxWhereMask(self, fn):
 
     self.assertAllClose(out_masked, out_filtered)
 
-  def testNormalizeWhereMask(self):
+  def testStandardizeWhereMask(self):
     x = jnp.array([5.5, 1.3, -4.2, 0.9])
     m = jnp.array([True, False, True, True])
     x_filtered = jnp.take(x, jnp.array([0, 2, 3]))
 
-    out_masked = jnp.take(nn.normalize(x, where=m), jnp.array([0, 2, 3]))
-    out_filtered = nn.normalize(x_filtered)
+    out_masked = jnp.take(nn.standardize(x, where=m), jnp.array([0, 2, 3]))
+    out_filtered = nn.standardize(x_filtered)
 
     self.assertAllClose(out_masked, out_filtered)