Update second_order package internal structure.

Use the new subpackage structure where individual files are private (hence their names being prefixed with `_`) and you can only import the entire subpackages (e.g. `from optax import second_order`, instead of `from optax.second_order import hessian`)

PiperOrigin-RevId: 581246429

optax/second_order/__init__.py

@@ -14,6 +14,6 @@
 # ==============================================================================
 """The second order optimisation sub-package."""
 
-from optax.second_order.fisher import fisher_diag
-from optax.second_order.hessian import hessian_diag
-from optax.second_order.hessian import hvp
+from optax.second_order._fisher import fisher_diag
+from optax.second_order._hessian import hessian_diag
+from optax.second_order._hessian import hvp

optax/second_order/fisher.py → optax/second_order/_fisher.py

@@ -25,15 +25,15 @@
 from jax import flatten_util
 import jax.numpy as jnp
 
-from optax.second_order import base
+from optax.second_order import _base
 
 
 def _ravel(p: Any) -> jax.Array:
   return flatten_util.ravel_pytree(p)[0]
 
 
 def fisher_diag(
-    negative_log_likelihood: base.LossFn,
+    negative_log_likelihood: _base.LossFn,
     params: Any,
     inputs: jax.Array,
     targets: jax.Array,

optax/second_order/hessian.py → optax/second_order/_hessian.py

@@ -25,15 +25,15 @@
 from jax import flatten_util
 import jax.numpy as jnp
 
-from optax.second_order import base
+from optax.second_order import _base
 
 
 def _ravel(p: Any) -> jax.Array:
   return flatten_util.ravel_pytree(p)[0]
 
 
 def hvp(
-    loss: base.LossFn,
+    loss: _base.LossFn,
     v: jax.Array,
     params: Any,
     inputs: jax.Array,
@@ -58,7 +58,7 @@ def hvp(
 
 
 def hessian_diag(
-    loss: base.LossFn,
+    loss: _base.LossFn,
     params: Any,
     inputs: jax.Array,
     targets: jax.Array,

optax/second_order/hessian_test.py → optax/second_order/_hessian_test.py

@@ -26,7 +26,7 @@
 import jax.numpy as jnp
 import numpy as np
 
-from optax import second_order
+from optax.second_order import _hessian
 
 
 NUM_CLASSES = 2
@@ -77,7 +77,7 @@ def jax_hessian_diag(loss_fun, params, inputs, targets):
   @chex.all_variants
   def test_hessian_diag(self):
     hessian_diag_fn = self.variant(
-        functools.partial(second_order.hessian_diag, self.loss_fn))
+        functools.partial(_hessian.hessian_diag, self.loss_fn))
     actual = hessian_diag_fn(self.parameters, self.data, self.labels)
     np.testing.assert_array_almost_equal(self.hessian, actual, 5)