%pip --quiet install objax

"""
Explanation: Objax Basics Tutorial
This tutorial introduces basic Objax concepts.
Prerequisites
Objax is a machine learning library written in Python which works on top of JAX. Readers should therefore have some familiarity with following:

Python. If you're new to Python or need a refresher check Python Tutorial or Python Beginner's Guide.
NumPy is a library for mathematical computations. Many of JAX primitives are based on NumPy and have the same syntax. NumPy is also useful for data manipulation outside of JAX/Objax. NumPy quickstart covers most of the needed topics. More information can be found on NumPy documentation site.
JAX can be described as NumPy with gradients which runs on accelerators (GPU and TPU). The JAX quickstart covers most of the concepts needed to understand Objax.

Installation and imports
Let's first install Objax:
End of explanation
"""


import jax.numpy as jn
import numpy as np
import objax

"""
Explanation: After Objax is installed, you can import all necessary modules:
End of explanation
"""


# Providing explicit values
jn.array([[1.0, 2.0, 3.0],
          [4.0, 5.0, 6.0]])

"""
Explanation: Tensors
Tensors are essentially multi-dimensional arrays. In JAX and Objax tensors can be placed on GPUs or TPUs to accelerate computations.
Objax relies on the jax.numpy.ndarray primitive from JAX to represent tensors. In turn, this primitive has a very similar API to NumPy ndarray.
Creating tensors
Tensors creation is very similar to NumPy and is done in multiple ways:

Provide explicit values to the tensor:
End of explanation
"""


arr = np.array([1.0, 2.0, 3.0])
jn.array(arr)

"""
Explanation: From a NumPy array:
End of explanation
"""


another_tensor = jn.array([[1.0, 2.0, 3.0],
                           [4.0, 5.0, 6.0]])
jn.array(another_tensor)

"""
Explanation: From another JAX tensor:
End of explanation
"""


jn.ones((3, 4))

jn.zeros((4, 5))

"""
Explanation: Using ones or zeros:
End of explanation
"""


t1 = jn.array([[1.0, 2.0, 3.0],
               [4.0, 5.0, 6.0]])
t2 = jn.ones(t1.shape) * 3
t1 + t2

"""
Explanation: As a result of a mathematical operation performed on other tensors:
End of explanation
"""


t = jn.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0]])
print('Number of dimensions: ', t.ndim)
print('Shape: ', t.shape)
print('Data type: ', t.dtype)

"""
Explanation: Tensor Properties
Similar to NumPy, one can explore various properties of tensors like shape, number of dimensions, or data type:
End of explanation
"""


np.array(t)

"""
Explanation: Converting tensors to numpy array
Objax/JAX tensors can be converted to NumPy arrays when needed to perform computations with NumPy:
End of explanation
"""


print('Original tensor t:\n', t)

try:
    t[0, 0] = -5.0  # This line will fail
except Exception as e:
    print(f'Exception {e}')

print('Tensor t after failed attempt to update:\n', t)

"""
Explanation: Tensors are immutable
One big difference between JAX ndarray and NumPy ndarray is that JAX ndarray is immutable:
End of explanation
"""


import jax.ops

print('Original tensor t:\n', t)
new_t = jax.ops.index_update(t, jax.ops.index[0, 0], -5.0)
print('Tensor t after update stays the same:\n', t)
print('Tensor new_t has updated value:\n', new_t)

"""
Explanation: Instead of updating an existing tensor, a new tensor should be created with updated elements. Updates of individual tensor elements is done using
index_update, index_add and some other JAX primitives:
End of explanation
"""


x = objax.random.normal((3, 4))
print(x)

"""
Explanation: More details about per-element updates of tensors can be found in JAX documentation.
In practice, most mathematical operations cover tensors as a whole while manual per-element updates are rarely needed.
Random numbers
It's very easy to generate random tensors in Objax:
End of explanation
"""


print('Random integers:', objax.random.randint((4,), low=0, high=10))
print('Random normal:', objax.random.normal((4,), mean=1.0, stddev=2.0))
print('Random truncated normal: ', objax.random.truncated_normal((4,), stddev=2.0))
print('Random uniform: ', objax.random.uniform((4,)))

"""
Explanation: There are multiple primitives for doing so:
End of explanation
"""


class SimpleModule(objax.Module):

    def __init__(self, length):
        self.v1 = objax.TrainVar(objax.random.normal((length,)))
        self.v2 = objax.TrainVar(jn.ones((2,)))

    def __call__(self, x):
        return jn.dot(x, self.v1)
        

m = SimpleModule(3)

"""
Explanation: Objax Variables and Modules
Objax Variables store values of tensors. Unlike tensors variables are mutable, i.e. the value which is stored in the variable can change.
Since tensors are immutable, variables change their value by replacing it with new tensors.
Variables are commonly used together with modules. The Module is a basic building block in Objax that stores variables and other modules. Also most modules are typically callable (i.e., implement the __call__ method) and when called perform some computations on their variables and sub-modules.
Here is an example of a simple module with one variable which performs the dot product of that variable with an input tensor:
End of explanation
"""


module_vars = m.vars()

print('type(module_vars): ', type(module_vars))
print('isinstance(module_vars, dict): ', isinstance(module_vars, dict))
print()

print('Variable names and shapes:')
print(module_vars)
print()

print('Variable names and values:')
for k, v in module_vars.items():
  print(f'{k}      {v.value}')

"""
Explanation: Modules keep track of all variables they own, including variables in sub-modules. The .vars() method list all the module's variables.
The method returns an instance of VarCollection which is a dictionary with several other useful methods.
End of explanation
"""


x = jn.ones((3,))
y = m(x)
print('Input: ', x)
print('Output: ', y)

"""
Explanation: If the __call__ method of the module takes tensors as input and outputs tensors then it can act as a mathematical function. In the general case __call__ can be a multivariate vector-values function.
The SimpleModule described above takes a vector of size length as input and outputs a scalar:
End of explanation
"""


x = jn.array([[1., 1., 1.],
              [1., 0., 0.],
              [0., 1., 0.],
              [0., 0., 1.]])
y = m(x)
print('Input:\n', x)
print('Output:\n', y)

"""
Explanation: The way jn.dot works allows us to run code on 2D tensors as well. In this case SimpleModule will treat the input as a batch of vectors, perform the dot product on each of them and return a vector with the results:
End of explanation
"""


print('Sequentially calling module on each row of 2D tensor:')
for idx in range(x.shape[0]):
  row_value = x[idx]
  out_value = m(row_value)
  print(f'm( {row_value} ) = {out_value}')

"""
Explanation: For comparison, here is the result of calling module m on each row of tensor x:
End of explanation
"""


def loss_fn(x):
    return m(x).sum()

print('loss_fn(x) = ', loss_fn(x))

"""
Explanation: How to compute gradients
As shown above, modules can act as mathematical functions. It's essential in machine learning to be able to compute gradients of functions and Objax provides a simple way to do this.
It's important to keep in mind that gradients are usually defined for scalar-value functions, while our modules can be vector-valued. In this case we need to define additional functions which will convert vector-valued output of the module into scalar. Then we can compute gradients of scalar valued function with respect to all input variables.
In the example with SimpleModule above let's define scalar-values loss function first:
End of explanation
"""


# Construct a module which computes gradients
gv = objax.GradValues(loss_fn, module_vars)

"""
Explanation: Then we create an objax.GradValues module which computes the gradients of loss_fn. We need to pass the function itself to the constructor of objax.GradValues as well as a VarCollection with the variables that loss_fn depends on:
End of explanation
"""


# gv returns both gradients and values of original function
grads, value = gv(x)

print('Gradients:')
for g, var_name in zip(grads, module_vars.keys()):
    print(g, ' w.r.t. ', var_name)
print()
print('Value: ', value)

"""
Explanation: gv is a module which returns the gradients of loss_fn and the values of loss_fn for the given input:
End of explanation
"""


jit_m = objax.Jit(m)
y = jit_m(x)
print('Input:\n', x)
print('Output:\n', y)

"""
Explanation: In the example above, grads is a list of gradients with respect to all variables from module_vars. The order of gradients in the grads list is the same as the order of corresponding variables in module_vars. So grads[0] is the gradient of the function w.r.t. m.v1 and grads[1] is the gradient w.r.t. m.v2.
Just-in-time compilation (JIT)
In the examples shown so far, the Python interpreter executes all operations one by one. This mode of execution becomes slow for larger and more complicated code.
Objax provides an easy and convenient way to compile sequence of operations using objax.Jit:
End of explanation
"""


def loss_fn(x, y):
  return ((m(x) - y) ** 2).sum()

jit_loss_fn = objax.Jit(loss_fn, module_vars)

x = objax.random.normal((2, 3))
y = jn.array((-1.0, 1.0))

print('x:\n ', x)
print('y:\n', y)

print('loss_fn(x, y): ', loss_fn(x, y))
print('jit_loss_fn(x, y): ', jit_loss_fn(x, y))

"""
Explanation: Objax.Jit can compile not only modules, but also functions and callables. In this case a variable collection should be passed to objax.Jit:
End of explanation
"""


x = objax.random.normal((100, 3))
# gv is a module define above which compute gradients
jit_gv = objax.Jit(gv)
print('Timing for jit_gv:')
%timeit jit_gv(x)
print('Timing for gv:')
%timeit gv(x)

"""
Explanation: There is no need to use JIT if you only need to compute a single JAX operation. However JIT can give significant speedups when multiple Objax/JAX operations are chained together. The next tutorial will show examples of how JIT is used in practice.
Nevertherless the difference in execution speed with and without JIT is evident even in this simple example:
End of explanation
"""


