0
# Backend Operations
1

2
Low-level operations and backend functionality for tensor operations, mathematical functions, and neural network primitives across different backend engines (JAX, TensorFlow, PyTorch, OpenVINO).
3

4
## Capabilities
5

6
### Core Tensor Operations
7

8
Basic tensor manipulation functions for creating, reshaping, and converting tensors.
9

10
```python { .api }
11
def cast(x, dtype):
12
    """
13
    Cast tensor to specified dtype.
14
    
15
    Parameters:
16
    - x: Input tensor
17
    - dtype: Target data type
18
    
19
    Returns:
20
    Tensor cast to target dtype
21
    """
22

23
def convert_to_tensor(x, dtype=None):
24
    """
25
    Convert input to backend tensor.
26
    
27
    Parameters:
28
    - x: Input data (array, list, etc.)
29
    - dtype: Target data type
30
    
31
    Returns:
32
    Backend tensor
33
    """
34

35
def convert_to_numpy(x):
36
    """
37
    Convert tensor to numpy array.
38
    
39
    Parameters:
40
    - x: Input tensor
41
    
42
    Returns:
43
    Numpy array
44
    """
45

46
def is_tensor(x):
47
    """
48
    Check if input is a tensor.
49
    
50
    Parameters:
51
    - x: Input to check
52
    
53
    Returns:
54
    Boolean indicating if input is tensor
55
    """
56

57
def shape(x):
58
    """
59
    Get shape of tensor.
60
    
61
    Parameters:
62
    - x: Input tensor
63
    
64
    Returns:
65
    Shape tuple
66
    """
67

68
def dtype(x):
69
    """
70
    Get data type of tensor.
71
    
72
    Parameters:
73
    - x: Input tensor
74
    
75
    Returns:
76
    Data type
77
    """
78

79
def reshape(x, new_shape):
80
    """
81
    Reshape tensor to new shape.
82
    
83
    Parameters:
84
    - x: Input tensor
85
    - new_shape: Target shape
86
    
87
    Returns:
88
    Reshaped tensor
89
    """
90

91
def transpose(x, axes=None):
92
    """
93
    Transpose tensor dimensions.
94
    
95
    Parameters:
96
    - x: Input tensor
97
    - axes: Permutation of dimensions (optional)
98
    
99
    Returns:
100
    Transposed tensor
101
    """
102

103
def expand_dims(x, axis):
104
    """
105
    Expand tensor dimensions.
106
    
107
    Parameters:
108
    - x: Input tensor
109
    - axis: Position to insert new axis
110
    
111
    Returns:
112
    Tensor with expanded dimensions
113
    """
114

115
def squeeze(x, axis=None):
116
    """
117
    Remove single-dimensional entries.
118
    
119
    Parameters:
120
    - x: Input tensor
121
    - axis: Specific axis to squeeze (optional)
122
    
123
    Returns:
124
    Squeezed tensor
125
    """
126
```
127

128
### Mathematical Operations
129

130
Core mathematical functions for tensor arithmetic and mathematical transformations.
131

132
```python { .api }
133
def add(x1, x2):
134
    """
135
    Element-wise addition.
136
    
137
    Parameters:
138
    - x1: First input tensor
139
    - x2: Second input tensor
140
    
141
    Returns:
142
    Element-wise sum
143
    """
144

145
def subtract(x1, x2):
146
    """Element-wise subtraction."""
147

148
def multiply(x1, x2):
149
    """Element-wise multiplication."""
150

151
def divide(x1, x2):
152
    """Element-wise division."""
153

154
def power(x1, x2):
155
    """Element-wise power."""
156

157
def sqrt(x):
158
    """
159
    Element-wise square root.
160
    
161
    Parameters:
162
    - x: Input tensor
163
    
164
    Returns:
165
    Square root of input
166
    """
167

168
def square(x):
169
    """Element-wise square."""
170

171
def abs(x):
172
    """Element-wise absolute value."""
173

174
def sign(x):
175
    """Element-wise sign function."""
176

177
def exp(x):
178
    """Element-wise exponential."""
179

180
def log(x):
181
    """Element-wise natural logarithm."""
182

183
def sin(x):
184
    """Element-wise sine."""
185

186
def cos(x):
187
    """Element-wise cosine."""
188

189
def tan(x):
190
    """Element-wise tangent."""
191

192
def sinh(x):
193
    """Element-wise hyperbolic sine."""
194

195
def cosh(x):
196
    """Element-wise hyperbolic cosine."""
197

198
def tanh(x):
199
    """Element-wise hyperbolic tangent."""
200

201
def ceil(x):
202
    """Element-wise ceiling."""
203

204
def floor(x):
205
    """Element-wise floor."""
206

207
def round(x):
208
    """Element-wise rounding."""
209

210
def maximum(x1, x2):
211
    """Element-wise maximum."""
212

213
def minimum(x1, x2):
214
    """Element-wise minimum."""
215

216
def clip(x, min_value, max_value):
217
    """
218
    Clip tensor values to range.
219
    
220
    Parameters:
221
    - x: Input tensor
222
    - min_value: Minimum value
223
    - max_value: Maximum value
224
    
225
    Returns:
226
    Clipped tensor
227
    """
228
```
229

230
### Linear Algebra Operations
231

232
Linear algebra functions for matrix operations and decompositions.
233

234
```python { .api }
235
def matmul(x1, x2):
236
    """
237
    Matrix multiplication.
238
    
239
    Parameters:
240
    - x1: First input tensor
241
    - x2: Second input tensor
242
    
243
    Returns:
244
    Matrix product
245
    """
246

247
def dot(x1, x2):
248
    """Dot product of two tensors."""
249

250
def tensordot(x1, x2, axes):
251
    """
252
    Tensor dot product along specified axes.
253
    
254
    Parameters:
255
    - x1: First input tensor
256
    - x2: Second input tensor
257
    - axes: Axes to contract over
258
    
259
    Returns:
260
    Tensor dot product
261
    """
262

263
def outer(x1, x2):
264
    """Outer product of two vectors."""
265

266
def inner(x1, x2):
267
    """Inner product of two tensors."""
268

269
def cross(x1, x2):
270
    """Cross product of two vectors."""
271

272
def norm(x, ord=None, axis=None, keepdims=False):
273
    """
274
    Compute tensor norm.
275
    
276
    Parameters:
277
    - x: Input tensor
278
    - ord: Order of norm
279
    - axis: Axis along which to compute norm
280
    - keepdims: Whether to keep dimensions
281
    
282
    Returns:
283
    Norm of tensor
284
    """
285

286
def det(x):
287
    """Determinant of square matrix."""
288

289
def inv(x):
290
    """Matrix inverse."""
291

292
def solve(a, b):
293
    """
294
    Solve linear system ax = b.
295
    
296
    Parameters:
297
    - a: Coefficient matrix
298
    - b: Right-hand side
299
    
300
    Returns:
301
    Solution x
302
    """
303

304
def eig(x):
305
    """Eigenvalues and eigenvectors."""
306

307
def svd(x, full_matrices=True):
308
    """
309
    Singular value decomposition.
310
    
311
    Parameters:
312
    - x: Input matrix
313
    - full_matrices: Whether to compute full-sized U and V
314
    
315
    Returns:
316
    Tuple (U, s, Vh) of SVD decomposition
317
    """
318

319
def qr(x):
320
    """QR decomposition."""
321

322
def cholesky(x):
323
    """Cholesky decomposition."""
324
```
325

326
### Array Creation and Manipulation
327

328
Functions for creating arrays and manipulating their structure.
329

330
```python { .api }
331
def zeros(shape, dtype='float32'):
332
    """
333
    Create tensor of zeros.
334
    
335
    Parameters:
336
    - shape: Shape of output tensor
337
    - dtype: Data type
338
    
339
    Returns:
340
    Tensor filled with zeros
341
    """
342

343
def ones(shape, dtype='float32'):
344
    """Create tensor of ones."""
345

346
def full(shape, fill_value, dtype='float32'):
347
    """
348
    Create tensor filled with value.
349
    
350
    Parameters:
351
    - shape: Shape of output tensor
352
    - fill_value: Value to fill with
353
    - dtype: Data type
354
    
355
    Returns:
356
    Tensor filled with specified value
357
    """
358

359
def eye(N, M=None, k=0, dtype='float32'):
360
    """
361
    Create identity matrix.
362
    
363
    Parameters:
364
    - N: Number of rows
365
    - M: Number of columns (defaults to N)
366
    - k: Index of diagonal
367
    - dtype: Data type
368
    
369
    Returns:
370
    Identity matrix
371
    """
372

373
def arange(start, stop=None, step=1, dtype='int32'):
374
    """
375
    Create range of values.
376
    
377
    Parameters:
378
    - start: Start value
379
    - stop: Stop value  
380
    - step: Step size
381
    - dtype: Data type
382
    
383
    Returns:
384
    Array of evenly spaced values
385
    """
386

387
def linspace(start, stop, num=50, endpoint=True, dtype='float32'):
388
    """
389
    Create linearly spaced values.
390
    
391
    Parameters:
392
    - start: Start value
393
    - stop: Stop value
394
    - num: Number of samples
395
    - endpoint: Whether to include endpoint
396
    - dtype: Data type
397
    
398
    Returns:
399
    Array of linearly spaced values
400
    """
401

402
def concatenate(tensors, axis=0):
403
    """
404
    Concatenate tensors along axis.
405
    
406
    Parameters:
407
    - tensors: List of tensors to concatenate
408
    - axis: Axis for concatenation
409
    
410
    Returns:
411
    Concatenated tensor
412
    """
413

414
def stack(tensors, axis=0):
415
    """
416
    Stack tensors along new axis.
417
    
418
    Parameters:
419
    - tensors: List of tensors to stack
420
    - axis: Axis for stacking
421
    
422
    Returns:
423
    Stacked tensor
424
    """
425

426
def split(x, indices_or_sections, axis=0):
427
    """
428
    Split tensor into sub-tensors.
429
    
430
    Parameters:
431
    - x: Input tensor
432
    - indices_or_sections: Split points or number of sections
433
    - axis: Axis along which to split
434
    
435
    Returns:
436
    List of sub-tensors
437
    """
438

439
def unstack(x, num=None, axis=0):
440
    """
441
    Unstack tensor along axis.
442
    
443
    Parameters:
444
    - x: Input tensor
445
    - num: Number of tensors to unstack
446
    - axis: Axis to unstack along
447
    
448
    Returns:
449
    List of unstacked tensors
450
    """
451
```
452

453
### Statistical Operations
454

455
Functions for computing statistics and aggregations over tensors.
456

457
```python { .api }
458
def mean(x, axis=None, keepdims=False):
459
    """
460
    Compute mean along axis.
461
    
462
    Parameters:
463
    - x: Input tensor
464
    - axis: Axis to compute mean over
465
    - keepdims: Whether to keep dimensions
466
    
467
    Returns:
468
    Mean of tensor
469
    """
470

471
def sum(x, axis=None, keepdims=False):
472
    """Compute sum along axis."""
473

474
def prod(x, axis=None, keepdims=False):
475
    """Compute product along axis."""
476

477
def max(x, axis=None, keepdims=False):
478
    """Compute maximum along axis."""
479

480
def min(x, axis=None, keepdims=False):
481
    """Compute minimum along axis."""
482

483
def std(x, axis=None, keepdims=False):
484
    """Compute standard deviation along axis."""
485

486
def var(x, axis=None, keepdims=False):
487
    """Compute variance along axis."""
488

489
def argmax(x, axis=None):
490
    """
491
    Indices of maximum values.
492
    
493
    Parameters:
494
    - x: Input tensor
495
    - axis: Axis to find argmax over
496
    
497
    Returns:
498
    Indices of maximum values
499
    """
500

501
def argmin(x, axis=None):
502
    """Indices of minimum values."""
503

504
def argsort(x, axis=-1):
505
    """
506
    Indices that would sort tensor.
507
    
508
    Parameters:
509
    - x: Input tensor
510
    - axis: Axis to sort along
511
    
512
    Returns:
513
    Indices that would sort tensor
514
    """
515

516
def sort(x, axis=-1):
517
    """Sort tensor along axis."""
518

519
def top_k(x, k, sorted=True):
520
    """
521
    Find top k values and indices.
522
    
523
    Parameters:
524
    - x: Input tensor
525
    - k: Number of top values to find
526
    - sorted: Whether to sort output
527
    
528
    Returns:
529
    Tuple of (values, indices)
530
    """
531
```
532

533
### Logical and Comparison Operations
534

535
Functions for logical operations and comparisons between tensors.
536

537
```python { .api }
538
def equal(x1, x2):
539
    """Element-wise equality."""
540

541
def not_equal(x1, x2):
542
    """Element-wise inequality."""
543

544
def less(x1, x2):
545
    """Element-wise less than."""
546

547
def less_equal(x1, x2):
548
    """Element-wise less than or equal."""
549

550
def greater(x1, x2):
551
    """Element-wise greater than."""
552

553
def greater_equal(x1, x2):
554
    """Element-wise greater than or equal."""
555

556
def logical_and(x1, x2):
557
    """Element-wise logical AND."""
558

559
def logical_or(x1, x2):
560
    """Element-wise logical OR."""
561

562
def logical_not(x):
563
    """Element-wise logical NOT."""
564

565
def logical_xor(x1, x2):
566
    """Element-wise logical XOR."""
567

568
def all(x, axis=None, keepdims=False):
569
    """
570
    Test whether all elements are True.
571
    
572
    Parameters:
573
    - x: Input tensor
574
    - axis: Axis to test over
575
    - keepdims: Whether to keep dimensions
576
    
577
    Returns:
578
    Boolean result
579
    """
580

581
def any(x, axis=None, keepdims=False):
582
    """Test whether any elements are True."""
583

584
def where(condition, x1, x2):
585
    """
586
    Select elements from x1 or x2 based on condition.
587
    
588
    Parameters:
589
    - condition: Boolean condition tensor
590
    - x1: Tensor to select from when condition is True
591
    - x2: Tensor to select from when condition is False
592
    
593
    Returns:
594
    Selected elements
595
    """
596
```
597

598
### Neural Network Operations
599

600
Specialized operations commonly used in neural networks.
601

602
```python { .api }
603
def relu(x):
604
    """
605
    Rectified Linear Unit activation.
606
    
607
    Parameters:
608
    - x: Input tensor
609
    
610
    Returns:
611
    ReLU activated tensor
612
    """
613

614
def relu6(x):
615
    """ReLU capped at 6."""
616

617
def elu(x, alpha=1.0):
618
    """Exponential Linear Unit."""
619

620
def selu(x):
621
    """Scaled Exponential Linear Unit."""
622

623
def gelu(x, approximate=False):
624
    """Gaussian Error Linear Unit."""
625

626
def sigmoid(x):
627
    """Sigmoid activation."""
628

629
def softmax(x, axis=-1):
630
    """
631
    Softmax activation.
632
    
633
    Parameters:
634
    - x: Input tensor
635
    - axis: Axis to apply softmax over
636
    
637
    Returns:
638
    Softmax activated tensor
639
    """
640

641
def log_softmax(x, axis=-1):
642
    """Log-softmax activation."""
643

644
def softplus(x):
645
    """Softplus activation."""
646

647
def swish(x):
648
    """Swish activation."""
649

650
def conv(inputs, kernel, strides=1, padding='valid', data_format=None, 
651
         dilation_rate=1):
652
    """
653
    N-dimensional convolution.
654
    
655
    Parameters:
656
    - inputs: Input tensor
657
    - kernel: Convolution kernel
658
    - strides: Stride of convolution
659
    - padding: Padding mode
660
    - data_format: Data format
661
    - dilation_rate: Dilation rate
662
    
663
    Returns:
664
    Convolution output
665
    """
666

667
def conv_transpose(inputs, kernel, strides, padding='valid', output_padding=None,
668
                   data_format=None, dilation_rate=1):
669
    """Transposed convolution (deconvolution)."""
670

671
def depthwise_conv(inputs, kernel, strides=1, padding='valid', data_format=None,
672
                   dilation_rate=1):
673
    """Depthwise convolution."""
674

675
def max_pool(inputs, pool_size, strides=None, padding='valid', data_format=None):
676
    """
677
    Max pooling operation.
678
    
679
    Parameters:
680
    - inputs: Input tensor
681
    - pool_size: Size of pooling window
682
    - strides: Stride of pooling
683
    - padding: Padding mode
684
    - data_format: Data format
685
    
686
    Returns:
687
    Max pooled output
688
    """
689

690
def average_pool(inputs, pool_size, strides=None, padding='valid', data_format=None):
691
    """Average pooling operation."""
692

693
def batch_normalization(x, mean, variance, offset=None, scale=None, 
694
                       variance_epsilon=1e-3):
695
    """
696
    Batch normalization.
697
    
698
    Parameters:
699
    - x: Input tensor
700
    - mean: Mean for normalization
701
    - variance: Variance for normalization
702
    - offset: Offset parameter (beta)
703
    - scale: Scale parameter (gamma)
704
    - variance_epsilon: Small constant for numerical stability
705
    
706
    Returns:
707
    Normalized tensor
708
    """
709

710
def layer_normalization(x, scale=None, offset=None, axis=-1, epsilon=1e-3):
711
    """Layer normalization."""
712

713
def dropout(x, rate, noise_shape=None, seed=None):
714
    """
715
    Dropout regularization.
716
    
717
    Parameters:
718
    - x: Input tensor
719
    - rate: Dropout rate
720
    - noise_shape: Shape of dropout mask
721
    - seed: Random seed
722
    
723
    Returns:
724
    Dropout applied tensor
725
    """
726
```
727

728
### Control Flow Operations
729

730
Operations for conditional execution and loops in computational graphs.
731

732
```python { .api }
733
def cond(pred, true_fn, false_fn):
734
    """
735
    Conditional execution.
736
    
737
    Parameters:
738
    - pred: Boolean predicate
739
    - true_fn: Function to execute if True
740
    - false_fn: Function to execute if False
741
    
742
    Returns:
743
    Result of executed function
744
    """
745

746
def switch(branch_index, branch_fns, default=None):
747
    """
748
    Switch between multiple functions.
749
    
750
    Parameters:
751
    - branch_index: Index of branch to execute
752
    - branch_fns: List of functions
753
    - default: Default function if index out of range
754
    
755
    Returns:
756
    Result of selected function
757
    """
758

759
def while_loop(cond, body, loop_vars, maximum_iterations=None):
760
    """
761
    While loop operation.
762
    
763
    Parameters:
764
    - cond: Loop condition function
765
    - body: Loop body function
766
    - loop_vars: Initial loop variables
767
    - maximum_iterations: Maximum number of iterations
768
    
769
    Returns:
770
    Final loop variables
771
    """
772

773
def fori_loop(lower, upper, body_fun, init_val):
774
    """
775
    For loop with index range.
776
    
777
    Parameters:
778
    - lower: Lower bound of range
779
    - upper: Upper bound of range  
780
    - body_fun: Function to execute in loop
781
    - init_val: Initial value
782
    
783
    Returns:
784
    Final accumulated value
785
    """
786

787
def scan(f, init, xs, length=None, reverse=False, unroll=1):
788
    """
789
    Scan operation (sequential application).
790
    
791
    Parameters:
792
    - f: Function to apply
793
    - init: Initial carry value
794
    - xs: Input sequence
795
    - length: Length of sequence
796
    - reverse: Whether to scan in reverse
797
    - unroll: Unroll factor for optimization
798
    
799
    Returns:
800
    Tuple of (final_carry, outputs)
801
    """
802
```
803

804
## Usage Examples
805

806
### Basic Tensor Operations
807

808
```python
809
import keras.ops as ops
810
import numpy as np
811

812
# Create tensors
813
x = ops.array([1, 2, 3, 4, 5])
814
y = ops.array([2, 3, 4, 5, 6])
815

816
# Basic arithmetic
817
sum_result = ops.add(x, y)
818
product = ops.multiply(x, y)
819
power_result = ops.power(x, 2)
820

821
# Mathematical functions
822
exp_result = ops.exp(x)
823
sin_result = ops.sin(x)
824
sqrt_result = ops.sqrt(x)
825

826
print(f"Sum: {ops.convert_to_numpy(sum_result)}")
827
print(f"Product: {ops.convert_to_numpy(product)}")
828
```
829

830
### Matrix Operations
831

832
```python
833
import keras.ops as ops
834

835
# Create matrices
836
A = ops.array([[1, 2], [3, 4]], dtype='float32')
837
B = ops.array([[5, 6], [7, 8]], dtype='float32')
838

839
# Matrix multiplication
840
C = ops.matmul(A, B)
841

842
# Linear algebra operations
843
det_A = ops.det(A)
844
inv_A = ops.inv(A)
845
norm_A = ops.norm(A)
846

847
print(f"Matrix product:\n{ops.convert_to_numpy(C)}")
848
print(f"Determinant: {ops.convert_to_numpy(det_A)}")
849
```
850

851
### Neural Network Operations
852

853
```python
854
import keras.ops as ops
855

856
# Activation functions
857
x = ops.array([-2, -1, 0, 1, 2], dtype='float32')
858

859
relu_out = ops.relu(x)
860
sigmoid_out = ops.sigmoid(x)
861
softmax_out = ops.softmax(ops.array([[1, 2, 3], [4, 5, 6]], dtype='float32'))
862

863
print(f"ReLU: {ops.convert_to_numpy(relu_out)}")
864
print(f"Sigmoid: {ops.convert_to_numpy(sigmoid_out)}")
865
print(f"Softmax:\n{ops.convert_to_numpy(softmax_out)}")
866

867
# Convolution operation
868
inputs = ops.random.normal((1, 28, 28, 1))  # Batch of 1 image
869
kernel = ops.random.normal((3, 3, 1, 32))   # 3x3 kernel, 32 filters
870

871
conv_out = ops.conv(inputs, kernel, strides=1, padding='same')
872
print(f"Convolution output shape: {ops.shape(conv_out)}")
873
```
874

875
### Statistical Operations
876

877
```python
878
import keras.ops as ops
879

880
# Create sample data
881
data = ops.random.normal((100, 10))  # 100 samples, 10 features
882

883
# Compute statistics
884
data_mean = ops.mean(data, axis=0)  # Mean along samples
885
data_std = ops.std(data, axis=0)    # Std along samples
886
data_max = ops.max(data, axis=1)    # Max along features
887
data_argmax = ops.argmax(data, axis=1)  # Index of max
888

889
print(f"Mean shape: {ops.shape(data_mean)}")
890
print(f"Std shape: {ops.shape(data_std)}")
891
print(f"Max shape: {ops.shape(data_max)}")
892
```
893

894
### Custom Operations with Control Flow
895

896
```python
897
import keras.ops as ops
898

899
def custom_activation(x, threshold=0.0):
900
    """Custom activation function using control flow."""
901
    return ops.where(
902
        ops.greater(x, threshold),
903
        x,  # Keep positive values
904
        ops.multiply(x, 0.1)  # Scale negative values
905
    )
906

907
def iterative_process(init_val, num_steps):
908
    """Example of iterative computation."""
909
    def step_fn(i, val):
910
        return val + ops.sin(ops.cast(i, 'float32'))
911
    
912
    return ops.fori_loop(0, num_steps, step_fn, init_val)
913

914
# Use custom operations
915
x = ops.array([-2, -1, 0, 1, 2], dtype='float32')
916
activated = custom_activation(x, threshold=0.5)
917

918
result = iterative_process(ops.array(0.0), 10)
919
print(f"Final result: {ops.convert_to_numpy(result)}")
920
```