Skip to content

[JAX FE] Correct creation of numpy dtypes #25278

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mryzhov merged 3 commits into from
Jun 28, 2024
Merged

[JAX FE] Correct creation of numpy dtypes #25278

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8ae55a2
[JAX FE]Correct creation of numpy dtypes
rkazants Jun 28, 2024
20c9558
Simplify logic of checking dtypes
rkazants Jun 28, 2024
0a12ceb
Correct labeling for JAX FE sources
rkazants Jun 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .github/labeler.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,7 @@
'category: TF FE':
- 'src/frontends/tensorflow/**/*'
- 'src/frontends/tensorflow_common/**/*'
- 'src/bindings/python/src/openvino/frontend/tensorflow/**/*'
- 'tests/layer_tests/tensorflow_tests/**/*'
- 'tests/layer_tests/tensorflow2_keras_tests/**/*'
- 'tests/layer_tests/jax_tests/**/*'
Expand All @@ -163,6 +164,7 @@

'category: JAX FE':
- 'src/frontends/jax/**/*'
- 'src/bindings/python/src/openvino/frontend/jax/**/*'
- 'tests/layer_tests/jax_tests/**/*'

'category: tools':
Expand Down
24 changes: 12 additions & 12 deletions src/bindings/python/src/openvino/frontend/jax/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,17 +11,17 @@
from openvino.runtime import op, Type as OVType, Shape, OVAny

numpy_to_ov_type_map = {
np.dtypes.Float32DType: OVType.f32,
np.dtypes.BoolDType: OVType.boolean,
np.float32: OVType.f32,
bool: OVType.boolean,
jax.dtypes.bfloat16: OVType.bf16, # TODO: check this
np.dtypes.Float16DType: OVType.f16,
np.dtypes.Float32DType: OVType.f32,
np.dtypes.Float64DType: OVType.f64,
np.dtypes.UInt8DType: OVType.u8,
np.dtypes.Int8DType: OVType.i8,
np.dtypes.Int16DType: OVType.i16,
np.dtypes.Int32DType: OVType.i32,
np.dtypes.Int64DType: OVType.i64,
np.float16: OVType.f16,
np.float32: OVType.f32,
np.float64: OVType.f64,
np.uint8: OVType.u8,
np.int8: OVType.i8,
np.int16: OVType.i16,
np.int32: OVType.i32,
np.int64: OVType.i64,
}

jax_to_ov_type_map = {
Expand Down Expand Up @@ -74,7 +74,7 @@ def get_ov_type_for_value(value):
if value.aval.dtype in jax_to_ov_type_map:
return OVAny(jax_to_ov_type_map[value.aval.dtype])
for k, v in numpy_to_ov_type_map.items():
if isinstance(value.aval.dtype, k):
if value.aval.dtype == k:
return OVAny(v)
for k, v in basic_to_ov_type_map.items():
if isinstance(value.aval.dtype, k):
Expand All @@ -88,7 +88,7 @@ def get_ov_type_from_jax_type(dtype):
if dtype in jax_to_ov_type_map:
return OVAny(jax_to_ov_type_map[dtype])
for k, v in numpy_to_ov_type_map.items():
if isinstance(dtype, k):
if dtype == k:
return OVAny(v)
for k, v in basic_to_ov_type_map.items():
if isinstance(dtype, k):
Expand Down