neural_network.optimizers.muon_optimizer
========================================

.. py:module:: neural_network.optimizers.muon_optimizer

.. autoapi-nested-parse::

   Muon Optimizer

   Implements Muon optimizer for neural network hidden layers using NumPy.
   Muon uses Newton-Schulz orthogonalization iterations for improved convergence.

   Reference: https://kellerjordan.github.io/posts/muon/
   Author: Adhithya Laxman Ravi Shankar Geetha
   Date: 2025.10.21



Attributes
----------

.. autoapisummary::

   neural_network.optimizers.muon_optimizer.optimizer


Classes
-------

.. autoapisummary::

   neural_network.optimizers.muon_optimizer.Muon


Module Contents
---------------

.. py:class:: Muon(learning_rate: float = 0.02, momentum: float = 0.95, ns_steps: int = 5)

   Muon optimizer for hidden layer weight matrices.

   Applies Newton-Schulz orthogonalization to gradients before updates.


   .. py:method:: newton_schulz_orthogonalize(matrix: numpy.ndarray) -> numpy.ndarray

      Orthogonalize matrix using Newton-Schulz iterations.

      Args:
          matrix (np.ndarray): Input matrix.

      Returns:
          np.ndarray: Orthogonalized matrix.

      >>> optimizer = Muon()
      >>> mat = np.array([[1.0, 0.5], [0.5, 1.0]])
      >>> orth = optimizer.newton_schulz_orthogonalize(mat)
      >>> orth.shape
      (2, 2)



   .. py:method:: update(param_id: int, params: numpy.ndarray, gradients: numpy.ndarray) -> numpy.ndarray

      Update parameters using Muon.

      Args:
          param_id (int): Unique identifier for parameter group.
          params (np.ndarray): Current parameters.
          gradients (np.ndarray): Gradients of parameters.

      Returns:
          np.ndarray: Updated parameters.

      >>> optimizer = Muon(learning_rate=0.1, momentum=0.9)
      >>> params = np.array([[1.0, 2.0], [3.0, 4.0]])
      >>> grads = np.array([[0.1, 0.2], [0.3, 0.4]])
      >>> updated = optimizer.update(0, params, grads)
      >>> updated.shape
      (2, 2)



   .. py:attribute:: learning_rate
      :value: 0.02



   .. py:attribute:: momentum
      :value: 0.95



   .. py:attribute:: ns_steps
      :value: 5



   .. py:attribute:: velocity
      :type:  dict[int, numpy.ndarray]


.. py:data:: optimizer