go-xla/compute/xla/capabilities.go at 15727bfdbecc532dab78cbb2839f17915df7ff53 · gomlx/go-xla · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
// Copyright 2023-2026 The GoMLX Authors. SPDX-License-Identifier: Apache-2.0

package xla

import (
	"github.com/gomlx/compute"
	"github.com/gomlx/compute/dtypes"
)

// Capabilities of the Go backends: the set of supported operations and data types.
var Capabilities = compute.Capabilities{
	Operations: map[compute.OpType]bool{
		// Graph inputs (leaf nodes)
		compute.OpTypeParameter: true,
		compute.OpTypeConstant:  true,

		// Standard unary operations:
		compute.OpTypeAbs:        true,
		compute.OpTypeBitCount:   true,
		compute.OpTypeBitwiseNot: true,
		compute.OpTypeCeil:       true,
		compute.OpTypeClz:        true,
		compute.OpTypeCos:        true,
		compute.OpTypeErf:        true,
		compute.OpTypeExp:        true,
		compute.OpTypeExpm1:      true,
		compute.OpTypeFloor:      true,
		compute.OpTypeIsFinite:   true,
		compute.OpTypeIsNaN:      true,
		compute.OpTypeLog1p:      true,
		compute.OpTypeLog:        true,
		compute.OpTypeLogicalNot: true,
		compute.OpTypeLogistic:   true,
		compute.OpTypeNeg:        true,
		compute.OpTypeRound:      true,
		compute.OpTypeRsqrt:      true,
		compute.OpTypeSign:       true,
		compute.OpTypeSin:        true,
		compute.OpTypeSqrt:       true,
		compute.OpTypeTanh:       true,

		// Standard binary operations:
		compute.OpTypeAdd:        true,
		compute.OpTypeBitwiseAnd: true,
		compute.OpTypeBitwiseOr:  true,
		compute.OpTypeBitwiseXor: true,
		compute.OpTypeDiv:        true,
		compute.OpTypeLogicalAnd: true,
		compute.OpTypeLogicalOr:  true,
		compute.OpTypeLogicalXor: true,
		compute.OpTypeMax:        true,
		compute.OpTypeMin:        true,
		compute.OpTypeMul:        true,
		compute.OpTypePow:        true,
		compute.OpTypeRem:        true,
		compute.OpTypeSub:        true,

		// Comparison operators.
		compute.OpTypeEqual:                    true,
		compute.OpTypeEqualTotalOrder:          true,
		compute.OpTypeGreaterOrEqual:           true,
		compute.OpTypeGreaterOrEqualTotalOrder: true,
		compute.OpTypeGreaterThan:              true,
		compute.OpTypeGreaterThanTotalOrder:    true,
		compute.OpTypeLessOrEqual:              true,
		compute.OpTypeLessOrEqualTotalOrder:    true,
		compute.OpTypeLessThan:                 true,
		compute.OpTypeLessThanTotalOrder:       true,
		compute.OpTypeNotEqual:                 true,
		compute.OpTypeNotEqualTotalOrder:       true,

		// Complex operations:
		compute.OpTypeComplex: true,
		compute.OpTypeConj:    true,
		compute.OpTypeImag:    true,
		compute.OpTypeReal:    true,

		// Other operations:
		compute.OpTypeArgMinMax:             true,
		compute.OpTypeBatchNormForInference: true,
		compute.OpTypeBatchNormForTraining:  true,
		compute.OpTypeBatchNormGradient:     true,
		compute.OpTypeBitcast:               true,
		compute.OpTypeBroadcastInDim:        true,
		compute.OpTypeClamp:                 true,
		compute.OpTypeConcatenate:           true,
		compute.OpTypeConvertDType:          true,
		compute.OpTypeConvGeneral:           true,
		compute.OpTypeDynamicSlice:          true,
		compute.OpTypeDynamicUpdateSlice:    true,
		compute.OpTypeDot:                   true,
		compute.OpTypeDotGeneral:            true,
		compute.OpTypeFFT:                   true,
		compute.OpTypeGather:                true,
		compute.OpTypeIdentity:              true,
		compute.OpTypeIota:                  true,
		compute.OpTypePad:                   true,
		compute.OpTypeReduceBitwiseAnd:      true,
		compute.OpTypeReduceBitwiseOr:       true,
		compute.OpTypeReduceBitwiseXor:      true,
		compute.OpTypeReduceLogicalAnd:      true,
		compute.OpTypeReduceLogicalOr:       true,
		compute.OpTypeReduceLogicalXor:      true,
		compute.OpTypeReduceMax:             true,
		compute.OpTypeReduceMin:             true,
		compute.OpTypeReduceProduct:         true,
		compute.OpTypeReduceSum:             true,
		compute.OpTypeReduceWindow:          true,
		compute.OpTypeReshape:               true,
		compute.OpTypeReverse:               true,
		compute.OpTypeRNGBitGenerator:       true,
		compute.OpTypeScatterSum:            true,
		compute.OpTypeScatterMax:            true,
		compute.OpTypeScatterMin:            true,
		compute.OpTypeSelectAndScatterMax:   true,
		compute.OpTypeSelectAndScatterMin:   true,
		compute.OpTypeShiftLeft:             true,
		compute.OpTypeShiftRightArithmetic:  true,
		compute.OpTypeShiftRightLogical:     true,
		compute.OpTypeSlice:                 true,
		compute.OpTypeTranspose:             true,
		compute.OpTypeWhere:                 true,
		compute.OpTypeOptimizationBarrier:   true,
		compute.OpTypeSchedulingBarrier:     true,

		// Collective (distributed across devices) operations:
		compute.OpTypeAllReduce: true,
	},

	Functions:   true,
	DynamicAxes: false,

	DTypes: map[dtypes.DType]bool{
		dtypes.Bool:       true,
		dtypes.Int8:       true,
		dtypes.Int16:      true,
		dtypes.Int32:      true,
		dtypes.Int64:      true,
		dtypes.Uint8:      true,
		dtypes.Uint16:     true,
		dtypes.Uint32:     true,
		dtypes.Uint64:     true,
		dtypes.Float32:    true,
		dtypes.Float64:    true,
		dtypes.BFloat16:   true,
		dtypes.Complex64:  true,
		dtypes.Complex128: true,
	},
}