Skip to content

Commit aabe9c6

Browse files
committed
Add F4E2M1FN type: HLO evaluator
1 parent 8e0305c commit aabe9c6

File tree

3 files changed

+24
-0
lines changed

3 files changed

+24
-0
lines changed

xla/hlo/evaluator/BUILD

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ cc_library(
3636
"hlo_evaluator_typed_visitor_int4.cc",
3737
"hlo_evaluator_typed_visitor_int64.cc",
3838
"hlo_evaluator_typed_visitor_int8.cc",
39+
"hlo_evaluator_typed_visitor_mxfloat.cc",
3940
"hlo_evaluator_typed_visitor_uint16.cc",
4041
"hlo_evaluator_typed_visitor_uint32.cc",
4142
"hlo_evaluator_typed_visitor_uint64.cc",

xla/hlo/evaluator/hlo_evaluator_typed_visitor.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1734,6 +1734,7 @@ extern template class HloEvaluatorTypedVisitor<double>;
17341734
extern template class HloEvaluatorTypedVisitor<complex64>;
17351735
extern template class HloEvaluatorTypedVisitor<complex128>;
17361736
extern template class HloEvaluatorTypedVisitor<bfloat16, float>;
1737+
extern template class HloEvaluatorTypedVisitor<tsl::float4_e2m1fn, float>;
17371738
extern template class HloEvaluatorTypedVisitor<tsl::float8_e5m2, float>;
17381739
extern template class HloEvaluatorTypedVisitor<tsl::float8_e4m3, float>;
17391740
extern template class HloEvaluatorTypedVisitor<tsl::float8_e4m3fn, float>;
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/* Copyright 2024 The OpenXLA Authors.
2+
3+
Licensed under the Apache License, Version 2.0 (the "License");
4+
you may not use this file except in compliance with the License.
5+
You may obtain a copy of the License at
6+
7+
http://www.apache.org/licenses/LICENSE-2.0
8+
9+
Unless required by applicable law or agreed to in writing, software
10+
distributed under the License is distributed on an "AS IS" BASIS,
11+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
See the License for the specific language governing permissions and
13+
limitations under the License.
14+
==============================================================================*/
15+
16+
#include "xla/hlo/evaluator/hlo_evaluator.h"
17+
#include "xla/hlo/evaluator/hlo_evaluator_typed_visitor.h"
18+
#include "tsl/platform/ml_dtypes.h"
19+
20+
namespace xla {
21+
template class HloEvaluatorTypedVisitor<tsl::float4_e2m1fn, float>;
22+
} // namespace xla

0 commit comments

Comments
 (0)