21. TensorFlow operators supported by the IPU

Supported operators for device: XLA_IPU_JIT

Operator

Type Constraint

Abs

T={float,half,int32,int64}

Acos

T={float,half,int32,int64}

Acosh

T={float,half}

Add

T={float,half,int32,int64}

AddN

T={float,half,int32,int64,variant}

AddV2

T={float,half,int32,int64}

AdjustContrastv2

T={float,half}

AdjustHue

T={float,half}

AdjustSaturation

T={float,half}

All

Tidx={int32,int64}

Any

Tidx={int32,int64}

ApproximateEqual

T={float,half,int32,int64}

ArgMax

output_type={int32,int64}
T={float,half,int32,int64}
Tidx={int32,int64}

ArgMin

output_type={int32,int64}
T={float,half,int32,int64}
Tidx={int32,int64}

Asin

T={float,half,int32,int64}

Asinh

T={float,half}

AssignAddVariableOp

dtype={float,half,int32,int64}

AssignSubVariableOp

dtype={float,half,int32,int64}

AssignVariableOp

dtype={bool,float,half,int32,int64}

Atan

T={float,half,int32,int64}

Atan2

T={float,half}

Atanh

T={float,half}

AvgPool

T={float,half}

AvgPool3D

T={float,half}

AvgPool3DGrad

T={float,half}

AvgPoolGrad

T={float,half}

BatchMatMul

T={float,half,int32,int64}

BatchMatMulV2

T={float,half,int32,int64}

BatchToSpace

Tidx={int32,int64}
T={bool,float,half,int32,int64}

BatchToSpaceND

Tcrops={int32,int64}
T={bool,float,half,int32,int64}
Tblock_shape={int32,int64}

BiasAdd

T={float,half,int32,int64}

BiasAddGrad

T={float,half,int32,int64}

BiasAddV1

T={float,half,int32,int64}

Bitcast

type={float,half,int32,int64}
T={float,half,int32,int64}

BitwiseAnd

T={int32,int64}

BitwiseOr

T={int32,int64}

BitwiseXor

T={int32,int64}

BroadcastArgs

T={int32,int64}

BroadcastGradientArgs

T={int32,int64}

BroadcastTo

Tidx={int32,int64}
T={bool,float,half,int32,int64,int8,uint8}

Bucketize

T={float,int32,int64}

Case

Tout={bool,float,half,int32,int64,resource,variant}
Tin={bool,float,half,int32,int64,resource,variant}

Cast

DstT={bool,float,half,int32,int64,int8,uint8}
SrcT={bool,float,half,int32,int64,int8,uint8}

Ceil

T={float,half}

Cholesky

T={float,half}

ClipByValue

T={float,half,int32,int64}

Concat

T={bool,float,half,int32,int64}

ConcatOffset

ConcatV2

Tidx={int32}
T={bool,float,half,int32,int64}

ConjugateTranspose

Tperm={int32,int64}
T={bool,float,half,int32,int64}

Const

dtype={bool,float,half,int32,int64,string}

ControlTrigger

Conv2D

T={float,half}

Conv2DBackpropFilter

T={float,half}

Conv2DBackpropInput

T={float,half}

Conv3D

T={float,half}

Conv3DBackpropFilterV2

T={float,half}

Conv3DBackpropInputV2

Tshape={int32,int64}
T={float,half}

Cos

T={float,half}

Cosh

T={float,half}

Cross

T={float,half,int32,int64}

Cumprod

Tidx={int32,int64}
T={float,half,int32}

Cumsum

Tidx={int32,int64}
T={float,half,int32}

DataFormatDimMap

T={int32,int64}

DataFormatVecPermute

T={int32,int64}

DepthToSpace

T={bool,float,half,int32,int64}

DepthwiseConv2dNative

T={float,half}

DepthwiseConv2dNativeBackpropFilter

T={float,half}

DepthwiseConv2dNativeBackpropInput

T={float,half}

Diag

T={float,half,int32,int64}

DiagPart

T={float,half,int32,int64}

Digamma

T={float,half}

Div

T={float,half,int32,int64}

DivNoNan

T={float,half}

DynamicStitch

T={bool,float,half,int32,int64}

Elu

T={float,half}

EluGrad

T={float,half}

Empty

dtype={bool,float,half,int32,int64}

EmptyTensorList

shape_type={int32,int64,variant}
element_dtype={bool,float,half,int32,int64,variant}

Equal

T={bool,float,half,int32,int64}

Erf

T={float,half}

Erfc

T={float,half}

Exp

T={float,half}

ExpandDims

Tdim={int32,int64}
T={bool,float,half,int32,int64,int8,uint8}

Expm1

T={float,half}

ExtractImagePatches

T={float,half,int32,int64}

FakeParam

dtype={bool,float,half,int32,int64}

FakeQuantWithMinMaxArgs

FakeQuantWithMinMaxArgsGradient

FakeQuantWithMinMaxVars

FakeQuantWithMinMaxVarsGradient

Fill

index_type={int32,int64}
T={bool,float,half,int32,int64}

Floor

T={float,half}

FloorDiv

T={float,half,int32,int64}

FloorMod

T={float,half,int32,int64}

FusedBatchNorm

T={float}

FusedBatchNormGrad

T={float}

FusedBatchNormGradV2

V={float,half}
T={float,half}
U={float,half}

FusedBatchNormGradV3

V={float,half}
T={float,half}
U={float,half}

FusedBatchNormV2

U={float,half}
T={float,half}

FusedBatchNormV3

U={float,half}
T={float,half}

Gather

Tindices={int32,int64}
Tparams={bool,float,half,int32,int64}

GatherNd

Tindices={int32,int64}
Tparams={bool,float,half,int32,int64}

GatherV2

Taxis={int32,int64}
Tparams={bool,float,half,int32,int64}
Tindices={int32,int64}

Greater

T={float,half,int32,int64}

GreaterEqual

T={float,half,int32,int64}

HSVToRGB

T={float,half}

IRFFT

IRFFT2D

IRFFT3D

Identity

T={bool,float,half,int32,int64,int8,resource,uint8,variant}

IdentityN

T={bool,float,half,int32,int64,resource,variant}

If

Tout={bool,float,half,int32,int64,resource,variant}
Tcond={bool,float,half,int32,int64,resource,variant}
Tin={bool,float,half,int32,int64,resource,variant}

InTopKV2

T={int32,int64}

Inv

T={float,half,int32,int64}

Invert

T={int32,int64}

InvertPermutation

T={int32}

IsFinite

T={float,half}

IsInf

T={float,half}

IsNan

T={float,half}

L2Loss

T={float,half}

LRN

T={float,half}

LRNGrad

T={float,half}

LeakyRelu

T={float,half}

LeakyReluGrad

T={float,half}

LeftShift

T={int32,int64}

Less

T={float,half,int32,int64}

LessEqual

T={float,half,int32,int64}

Lgamma

T={float,half}

LinSpace

Tidx={int32,int64}
T={float}

ListDiff

out_idx={int32,int64}
T={int32,int64}

Log

T={float,half}

Log1p

T={float,half}

LogSoftmax

T={float,half}

LogicalAnd

LogicalNot

LogicalOr

MatMul

T={float,half}

MatrixBandPart

Tindex={int32,int64}
T={bool,float,half,int32,int64}

MatrixDiag

T={bool,float,half,int32,int64}

MatrixDiagPart

T={bool,float,half,int32,int64}

MatrixDiagPartV2

T={bool,float,half,int32,int64}

MatrixDiagV2

T={bool,float,half,int32,int64}

MatrixInverse

T={float,half}

MatrixSetDiag

T={bool,float,half,int32,int64}

MatrixSetDiagV2

T={bool,float,half,int32,int64}

MatrixTriangularSolve

T={float,half}

Max

Tidx={int32,int64}
T={float,half,int32,int64}

MaxPool

T={float,half,int32,int64}

MaxPool3D

T={float,half}

MaxPool3DGrad

TInput={float,half}
T={float,half}

MaxPoolGrad

T={float,half,int32,int64}

MaxPoolGradGradV2

T={float}

MaxPoolGradV2

T={float,half,int32,int64}

MaxPoolV2

T={float,half,int32,int64}

Maximum

T={float,half,int32,int64}

Mean

Tidx={int32,int64}
T={float,half,int32,int64}

Min

Tidx={int32,int64}
T={float,half,int32,int64}

Minimum

T={float,half,int32,int64}

MirrorPad

Tpaddings={int32,int64}
T={bool,float,half,int32,int64}

Mod

T={float,half,int32,int64}

Mul

T={float,half,int32,int64}

MulNoNan

T={float,half}

Multinomial

output_dtype={int32,int64}
T={float,half,int32,int64}

Neg

T={float,half,int32,int64}

NextAfter

T={float}

NoOp

NotEqual

T={bool,float,half,int32,int64}

OneHot

TI={int32,int64}
T={bool,float,half,int32,int64}

OnesLike

T={bool,float,half,int32,int64}

Pack

T={bool,float,half,int32,int64}

Pad

Tpaddings={int32,int64}
T={bool,float,half,int32,int64}

PadV2

Tpaddings={int32,int64}
T={bool,float,half,int32,int64}

ParallelDynamicStitch

T={bool,float,half,int32,int64}

ParameterizedTruncatedNormal

T={int32,int64}
dtype={float}

PartitionedCall

Tout={bool,float,half,int32,int64,resource,string,variant}
Tin={bool,float,half,int32,int64,resource,string,variant}

PlaceholderWithDefault

dtype={bool,float,half,int32,int64}

Pow

T={float,half,int32,int64}

PreventGradient

T={bool,float,half,int32,int64}

Prod

Tidx={int32,int64}
T={float,half,int32,int64}

QuantizeAndDequantizeV2

T={float,half}

QuantizeAndDequantizeV3

T={float,half}

RFFT

RFFT2D

RFFT3D

RGBToHSV

T={float,half}

RandomShuffle

T={bool,float,half,int32,int64}

RandomStandardNormal

T={int32,int64}
dtype={float,half}

RandomUniform

T={int32,int64}
dtype={float,half}

RandomUniformInt

T={int32,int64}
Tout={int32,int64}

Range

Tidx={float,int32,int64}

Rank

T={bool,float,half,int32,int64}

ReadVariableOp

dtype={bool,float,half,int32,int64}

RealDiv

T={float,half,int32,int64}

Reciprocal

T={float,half,int32,int64}

ReciprocalGrad

T={float,half}

Relu

T={float,half,int32,int64}

Relu6

T={float,half,int32,int64}

Relu6Grad

T={float,half,int32,int64}

ReluGrad

T={float,half,int32,int64}

Reshape

Tshape={int32,int64}
T={bool,float,half,int32,int64,int8,uint8}

ResizeBilinear

T={float,half,int32,int64}

ResizeBilinearGrad

T={float,half}

ResizeNearestNeighbor

T={float,half,int32,int64}

ResizeNearestNeighborGrad

T={float,half,int32}

ResourceApplyAdaMax

T={float,half}

ResourceApplyAdadelta

T={float,half}

ResourceApplyAdagrad

T={float,half}

ResourceApplyAdagradDA

T={float,half}

ResourceApplyAdagradV2

T={float,half}

ResourceApplyAdam

T={float,half}

ResourceApplyAddSign

T={float,half}

ResourceApplyCenteredRMSProp

T={float,half}

ResourceApplyFtrl

T={float,half}

ResourceApplyFtrlV2

T={float,half}

ResourceApplyGradientDescent

T={float,half}

ResourceApplyKerasMomentum

T={float,half}

ResourceApplyMomentum

T={float,half}

ResourceApplyPowerSign

T={float,half}

ResourceApplyProximalAdagrad

T={float,half}

ResourceApplyProximalGradientDescent

T={float,half}

ResourceApplyRMSProp

T={float,half}

ResourceGather

Tindices={int32,int64}
dtype={bool,float,half,int32,int64}

ResourceScatterAdd

Tindices={int32,int64}
dtype={float,half,int32,int64}

ResourceScatterDiv

Tindices={int32,int64}
dtype={float,half,int32,int64}

ResourceScatterMax

Tindices={int32,int64}
dtype={float,half,int32,int64}

ResourceScatterMin

Tindices={int32,int64}
dtype={float,half,int32,int64}

ResourceScatterMul

Tindices={int32,int64}
dtype={float,half,int32,int64}

ResourceScatterNdAdd

Tindices={int32,int64}
T={bool,float,half,int32,int64}

ResourceScatterNdSub

Tindices={int32,int64}
T={bool,float,half,int32,int64}

ResourceScatterNdUpdate

Tindices={int32,int64}
T={bool,float,half,int32,int64}

ResourceScatterSub

Tindices={int32,int64}
dtype={float,half,int32,int64}

ResourceScatterUpdate

Tindices={int32,int64}
dtype={bool,float,half,int32,int64}

ResourceStridedSliceAssign

Index={int32,int64}
T={bool,float,half,int32,int64}

Reverse

T={bool,float,half,int32,int64}

ReverseSequence

Tlen={int32,int64}
T={bool,float,half,int32,int64}

ReverseV2

T={bool,float,half,int32,int64}
Tidx={int32,int64}

RightShift

T={int32,int64}

Rint

T={float,half}

Roll

Taxis={int32,int64}
T={bool,float,half,int32,int64}
Tshift={int32,int64}

Round

T={float,half,int32,int64}

Rsqrt

T={float,half}

RsqrtGrad

T={float,half}

ScatterNd

Tindices={int32,int64}
T={bool,float,half,int32,int64}

Select

T={bool,float,half,int32,int64}

SelectV2

T={bool,float,half,int32,int64}

SelfAdjointEigV2

T={float,half}

Selu

T={float,half}

SeluGrad

T={float,half}

Shape

out_type={int32,int64}
T={bool,float,half,int32,int64}

ShapeN

out_type={int32,int64}
T={bool,float,half,int32,int64}

Sigmoid

T={float,half}

SigmoidGrad

T={float,half}

Sign

T={float,half,int32,int64}

Sin

T={float,half}

Sinh

T={float,half}

Size

out_type={int32,int64}
T={bool,float,half,int32,int64}

Slice

Index={int32,int64}
T={bool,float,half,int32,int64}

Snapshot

T={bool,float,half,int32,int64}

Softmax

T={float,half}

SoftmaxCrossEntropyWithLogits

T={float,half}

Softplus

T={float,half}

SoftplusGrad

T={float,half}

Softsign

T={float,half}

SoftsignGrad

T={float,half}

SpaceToBatch

Tpaddings={int32,int64}
T={bool,float,half,int32,int64}

SpaceToBatchND

Tpaddings={int32,int64}
T={bool,float,half,int32,int64}
Tblock_shape={int32,int64}

SpaceToDepth

T={bool,float,half,int32,int64}

SparseMatMul

Tb={float}
Ta={float}

SparseSoftmaxCrossEntropyWithLogits

Tlabels={int32,int64}
T={float,half}

SparseToDense

Tindices={int32,int64}
T={bool,float,half,int32,int64}

Split

T={bool,float,half,int32,int64}

SplitV

Tlen={int32,int64}
T={bool,float,half,int32,int64}

Sqrt

T={float,half}

SqrtGrad

T={float,half}

Square

T={float,half,int32,int64}

SquaredDifference

T={float,half,int32,int64}

Squeeze

T={bool,float,half,int32,int64}

StackCloseV2

StackPopV2

elem_type={bool,float,half,int32,int64}

StackPushV2

T={bool,float,half,int32,int64}

StackV2

elem_type={bool,float,half,int32,int64}

StatefulPartitionedCall

Tout={bool,float,half,int32,int64,resource,string,variant}
Tin={bool,float,half,int32,int64,resource,string,variant}

StatefulStandardNormalV2

shape_dtype={bool,float,half,int32,int64}
dtype={float}

StatefulTruncatedNormal

shape_dtype={bool,float,half,int32,int64}
dtype={float}

StatefulUniform

shape_dtype={bool,float,half,int32,int64}
dtype={float}

StatefulUniformFullInt

shape_dtype={bool,float,half,int32,int64}
dtype={int32,int64}

StatefulUniformInt

shape_dtype={bool,float,half,int32,int64}
dtype={int32,int64}

StatelessIf

Tout={bool,float,half,int32,int64,resource,variant}
Tcond={bool,float,half,int32,int64,resource,variant}
Tin={bool,float,half,int32,int64,resource,variant}

StatelessMultinomial

T={float}
output_dtype={int32,int64}
Tseed={int32}

StatelessRandomNormal

Tseed={int32}
dtype={float}
T={int32,int64}

StatelessRandomUniform

Tseed={int32}
dtype={float}
T={int32,int64}

StatelessRandomUniformInt

Tseed={int32}
dtype={int32,int64}
T={int32,int64}

StatelessTruncatedNormal

Tseed={int32}
dtype={float}
T={int32,int64}

StatelessWhile

T={bool,float,half,int32,int64,resource,variant}

StopGradient

T={bool,float,half,int32,int64}

StridedSlice

Index={int32,int64}
T={bool,float,half,int32,int64}

StridedSliceGrad

Index={int32,int64}
T={bool,float,half,int32,int64}

Sub

T={float,half,int32,int64}

Sum

Tidx={int32,int64}
T={float,half,int32,int64}

Svd

T={float,half}

SymbolicGradient

Tout={bool,float,half,int32,int64}
Tin={bool,float,half,int32,int64}

Tan

T={float,half,int32,int64}

Tanh

T={float,half}

TanhGrad

T={float,half}

TensorArrayCloseV3

TensorArrayConcatV3

dtype={bool,float,half,int32,int64}

TensorArrayGatherV3

dtype={bool,float,half,int32,int64}

TensorArrayGradV3

TensorArrayReadV3

dtype={bool,float,half,int32,int64}

TensorArrayScatterV3

T={bool,float,half,int32,int64}

TensorArraySizeV3

TensorArraySplitV3

T={bool,float,half,int32,int64}

TensorArrayV3

dtype={bool,float,half,int32,int64}

TensorArrayWriteV3

T={bool,float,half,int32,int64}

TensorListElementShape

shape_type={int32,int64}

TensorListFromTensor

shape_type={int32,int64}
element_dtype={bool,float,half,int32,int64}

TensorListGather

element_dtype={bool,float,half,int32,int64}

TensorListGetItem

element_dtype={bool,float,half,int32,int64}

TensorListLength

TensorListPopBack

element_dtype={bool,float,half,int32,int64,variant}

TensorListPushBack

element_dtype={bool,float,half,int32,int64,variant}

TensorListReserve

shape_type={int32,int64}
element_dtype={bool,float,half,int32,int64}

TensorListSetItem

element_dtype={bool,float,half,int32,int64}

TensorListStack

element_dtype={bool,float,half,int32,int64}

TensorScatterAdd

Tindices={int32,int64}
T={bool,float,half,int32,int64}

TensorScatterSub

Tindices={int32,int64}
T={bool,float,half,int32,int64}

TensorScatterUpdate

Tindices={int32,int64}
T={bool,float,half,int32,int64}

Tile

Tmultiples={int32,int64}
T={bool,float,half,int32,int64}

TopKV2

T={float,int32}

Transpose

Tperm={int32,int64}
T={bool,float,half,int32,int64}

TruncateDiv

T={float,half,int32,int64}

TruncateMod

T={float,half,int32,int64}

TruncatedNormal

T={int32,int64}
dtype={float}

Unpack

T={bool,float,half,int32,int64}

UnsortedSegmentMax

Tnumsegments={int32,int64}
T={float,half,int32,int64}
Tindices={int32,int64}

UnsortedSegmentMin

Tnumsegments={int32,int64}
T={float,half,int32,int64}
Tindices={int32,int64}

UnsortedSegmentProd

Tnumsegments={int32,int64}
T={float,half,int32,int64}
Tindices={int32,int64}

UnsortedSegmentSum

Tnumsegments={int32,int64}
T={float,half,int32,int64}
Tindices={int32,int64}

VarIsInitializedOp

VariableShape

out_type={int32,int64}

While

T={bool,float,half,int32,int64,resource,variant}

Xdivy

T={float,half}

XlaBroadcastHelper

Tindices={int32,int64}
T={float,half,int32,int64}

XlaConv

Tindices={int32,int64}
T={float,half,int32,int64}

XlaDequantize

XlaDot

T={float,half,int32,int64}

XlaDynamicSlice

Tindices={int32,int64}
T={bool,float,half,int32,int64}

XlaDynamicUpdateSlice

Tindices={int32,int64}
T={bool,float,half,int32,int64}

XlaEinsum

T={float}

XlaIf

Tout={bool,float,half,int32,int64,resource,variant}
Tcond={bool,float,half,int32,int64,resource,variant}
Tin={bool,float,half,int32,int64,resource,variant}

XlaKeyValueSort

V={bool,float,half,int32,int64}
K={float,half,int32,int64}

XlaPad

Tindices={int32,int64}
T={bool,float,half,int32,int64}

XlaRecv

dtype={bool,float,half,int32,int64}

XlaReduce

T={float,half,int32,int64}

XlaReduceWindow

Tindices={int32,int64}
T={float,half,int32,int64}

XlaReplicaId

XlaSelectAndScatter

Tindices={int32,int64}
T={float,half,int32,int64}

XlaSelfAdjointEig

T={float,half}

XlaSend

T={bool,float,half,int32,int64}

XlaSort

T={bool,float,half,int32,int64}

XlaSvd

T={float,half}

XlaWhile

T={bool,float,half,int32,int64,resource,variant}

Xlogy

T={float,half}

ZerosLike

T={bool,float,half,int32,int64,variant}

_Arg

T={bool,float,half,int32,int64,int8,resource,uint8,variant}

_ArrayToList

out_types={bool,float,half,int32,int64}
T={bool,float,half,int32,int64}

_FusedBatchNormEx

U={float,half}
T={float,half}

_ListToArray

T={bool,float,half,int32,int64}
Tin={bool,float,half,int32,int64}

_Retval

T={bool,float,half,int32,int64,resource,variant}

_UnaryOpsComposition

T={float,half}

To regenerate this table, run:

bazel run -c opt -- tensorflow/compiler/tf2xla:tf2xla_supported_ops --device=XLA_IPU_JIT