try bfloat16 from tensorflow

Xilinx · Sep 16, 2024 · 3a3acde · 3a3acde
1 parent f2fe13c
commit 3a3acde
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 9 deletions.
diff --git a/programming_examples/basic/matrix_multiplication/single_core/aie2.py b/programming_examples/basic/matrix_multiplication/single_core/aie2.py
@@ -6,10 +6,12 @@
 # (c) Copyright 2023 AMD Inc.
 
 import argparse
-from bfloat16 import bfloat16
 import numpy as np
+import tensorflow as tf
 import sys
 
+bfloat16 = tf.bfloat16.as_numpy_dtype
+
 from aie.dialects.scf import yield_, for_ as range_
 from aie.dialects.aiex import npu_dma_memcpy_nd, npu_sync
 

diff --git a/python/extras/util.py b/python/extras/util.py
@@ -1,5 +1,4 @@
 from collections import defaultdict
-from bfloat16 import bfloat16
 import contextlib
 import ctypes
 import inspect
@@ -12,8 +11,11 @@
 from pathlib import Path
 from typing import Callable, List, Optional, Sequence, Tuple, Union
 
+import tensorflow as tf
 import numpy as np
 
+bfloat16 = tf.bfloat16.as_numpy_dtype
+
 from .meta import op_region_builder
 from ..extras import types as T
 from ..ir import (

diff --git a/python/requirements.txt b/python/requirements.txt
@@ -4,11 +4,11 @@ cmake==3.27.9
 dataclasses>=0.6, <=0.8
 filelock==3.13.1
 lit
-numpy<2
+numpy>=1.19.5, <=1.26
 pandas
 psutil
 pybind11>=2.9.0, <=2.10.3
 rich
 setuptools
-bfloat16
 wheel
+tensorflow-cpu
diff --git a/python/requirements_ml.txt b/python/requirements_ml.txt
@@ -1,4 +1,2 @@
 --index-url https://download.pytorch.org/whl/cpu
-numpy<2
 torch
-bfloat16
diff --git a/utils/mlir_aie_wheels/requirements.txt b/utils/mlir_aie_wheels/requirements.txt
@@ -2,10 +2,9 @@ cmake==3.27.9
 dataclasses
 importlib-metadata
 ninja
-numpy<2
+numpy>=1.19.5, <=1.26
 pip
 pybind11[global]>=2.10.4
 rich
 setuptools>=42
-wheel
-bfloat16
+wheel