Add new dtypes support for GEMM

[linamy85]

Change

• Make sure bfloat16, float32, float16, float8, float4 is supported for GEMM.
• Updated gemm_multiple_run* yaml files to include test cases.

Test

apiVersion: v1
kind: Pod
metadata:
  name: microbenchmark
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu7x
    cloud.google.com/gke-tpu-topology: 2x2x1
  containers:
  - name: tpu-job
    image: python:3.12
    ports:
    - containerPort: 8431
    securityContext:
      privileged: false
    command:
    - bash
    - -c
    - |
      set -ex

      git clone https://github.com/linamy85/accelerator-microbenchmarks.git
      cd accelerator-microbenchmarks
      git checkout feature/new-dtype
      pip install -r requirements.txt

      python3 Ironwood/src/run_benchmark.py --config=Ironwood/configs/training/gemm_multiple_run.yaml

      sleep 36000

    resources:
      requests:
        google.com/tpu: 4
      limits:
        google.com/tpu: 4

[chishuen]

LGTM

[junjieqian]

Thanks! Merged this one as doing tests.