Check Identity Matrix in Python

from math import fma


def main():
	"""
	Prints the product of two matrices.
	"""

	# Get the dimensions of matrices A and B
	rows_a = int(input("Enter row length of A: "))
	cols_a = int(input("Enter col length of A: "))
	cols_b = int(input("Enter col length of B: "))

	# Read the matrices from the user
	print("Enter matrix A:")
	a = read_matrix(rows_a, cols_a)
	print("Enter matrix B:")
	b = read_matrix(cols_a, cols_b)

	# Multiply the matrices together
	c = multiply(a, b)

	# Print the result
	print("The product is:")
	print_matrix(c)


def multiply(x: list[list[int]], y: list[list[int]]):
	"""
	Returns the product of two matrices.
	"""

	# Get the dimensions of the matrices
	rows_x, cols_x = len(x), len(x[0])
	rows_y, cols_y = len(y), len(y[0])

	# Check if the matrices can be multiplied
	if cols_x != rows_y:
		raise ValueError(
			"Total columns in 1st matrix must be equal to total rows in 2nd matrix."
		)

	# Create the result matrix
	z = [[0] * cols_y for _ in range(rows_x)]

	# Multiply the matrices
	for r in range(rows_x):
		for c in range(cols_y):
			for i in range(cols_x):
				z[r][c] = int(fma(x[r][i], y[i][c], z[r][c]))

	return z


def read_matrix(row_count: int, col_count: int):
	"""
	Reads a matrix from the console.
	"""

	matrix: list[list[int]] = []
	for _ in range(row_count):
		row = list(map(int, input().split()[:col_count]))
		matrix.append(row)

	return matrix


def print_matrix(matrix: list[list[int]]):
	"""
	Prints a matrix to the console.
	"""

	for row in matrix:
		for element in row:
			print(element, end=" ")
		print()


if __name__ == "__main__":
	main()