how to use JAX library for matrix manipulation

• SPMD
  
  JAX supports the Single-Program Multi-Data (SPMD) paradigm. This allows the same computation to be performed across multiple devices. JAX doesn't require that all data is partitioned across the same devices. This means data can be distributed over different devices such as GPUs, CPUs, and TPUs for parallel processing.

• jax.jit
  
  The jax.jit function improves the computational efficiency by applying Just-In-Time (JIT) compilation to a JAX function. The effectiveness of jax.jit stems from a fundamental principle that each function is decomposed into a sequence of primitive operations, where every element represents a fundamental unit of computation.
  

import jax.numpy as jnp
from jax import jit

# Define matrices
a = jnp.array([[1, 2], [3, 4]])
b = jnp.array([[5, 6], [7, 8]])

print("Matrix A:")
print(a)
print("\nMatrix B:")
print(b)

# Matrix addition
matrix_sum = a + b
print("\nMatrix Addition (A + B):")
print(matrix_sum)

# Element-wise matrix multiplication
elementwise = a * b
print("\nElement-wise Multiplication (A * B):")
print(elementwise)

# Matrix dot product
dot_product = jnp.dot(a, b)
# Or use the @ operator
dot_product_alt = a @ b
print("\nMatrix Dot Product (A @ B):")
print(dot_product)

# Matrix transposition
transpose = a.T
print("\nMatrix A Transpose:")
print(transpose)

# Inverse matrix
try:
    inv = jnp.linalg.inv(a)
    print("\nInverse of Matrix A:")
    print(inv)
except:
    print("\nMatrix A is not invertible")

# Determinant
det = jnp.linalg.det(a)
print("\nDeterminant of Matrix A:")
print(det)

# SPMD and JIT compilation example
@jit
def matrix_operations(x, y):
    """Matrix operation function using JIT compilation"""
    dot = x @ y
    transpose = x.T
    return dot, transpose

# Using the JIT compiled function
result_dot, result_transpose = matrix_operations(a, b)
print("\nJIT Compiled Results:")
print("Dot Product:")
print(result_dot)
print("Transpose:")
print(result_transpose)