Basic Matrix#

Basic Matrix Operations with JAX#

import jax as J
import jax.numpy as jnp

Addition and Subtraction#

A = jnp.array(
    [
        [4, 5, 6, 8, 9],
        [3, 5, 7, 8, 3],
        [1, 4, 6, 0, 2]
    ]
)

A

Array([[4, 5, 6, 8, 9],
       [3, 5, 7, 8, 3],
       [1, 4, 6, 0, 2]], dtype=int32)

B = jnp.array(
    [
        [10, 1, 6, 8, 9],
        [3, 12, 7, 8, 3],
        [1, 4, 6, -1, 2]
    ]
)

B

Array([[10,  1,  6,  8,  9],
       [ 3, 12,  7,  8,  3],
       [ 1,  4,  6, -1,  2]], dtype=int32)

# addition

A_plus_B = jnp.add(A, B)
A_plus_B

Array([[14,  6, 12, 16, 18],
       [ 6, 17, 14, 16,  6],
       [ 2,  8, 12, -1,  4]], dtype=int32)

# subtraction

A_minus_B = jnp.subtract(A, B)
A_minus_B

Array([[-6,  4,  0,  0,  0],
       [ 0, -7,  0,  0,  0],
       [ 0,  0,  0,  1,  0]], dtype=int32)

Multiplication#

# check shapes first!
print(f"A : {A.shape}  B : {B.shape}")

A : (3, 5)  B : (3, 5)

The dimensions of A and B don’t satisfy the rules of Matrix multiplication. 3X5 does not multiply with 3X5. So let’s create a new matrix here with 5X3 dimension.

C = jnp.array(
    [
        [1, 2, 3],
        [2, 3, 5],
        [3, 4, 6],
        [4, 5, 7],
        [5, 6, 8],
    ]
)

C

Array([[1, 2, 3],
       [2, 3, 5],
       [3, 4, 6],
       [4, 5, 7],
       [5, 6, 8]], dtype=int32)

A_mult_C = jnp.dot(A, C)
A_mult_C

Array([[109, 141, 201],
       [ 81, 107, 156],
       [ 37,  50,  75]], dtype=int32)

Inverting#

Erm, we need a invertible matrix here. You can’t just invert every matrix out there! An invertible matrix has two properties:

  1. It’s a square matrix (equal number of rows and columns)

  2. The determinant of the matrix is NOT 0

What is an invertible matrix

Luckily A_mult_C is a square matrix. Let’s see if this can be inverted.

D = A_mult_C

determinant = jnp.linalg.det(D)

if determinant != 0:
    print("Invertible")
else:
    print("Not Invertible")

Invertible

# time to invert
inverse = jnp.linalg.inv(D)
inverse

Array([[ 2.4193413 , -5.645129  ,  5.258035  ],
       [-3.2580476 ,  7.9354434 , -7.7741566 ],
       [ 0.97848994, -2.5053654 ,  2.6021404 ]], dtype=float32)