class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=False):
        super().__init__()
        # Define the convolutional, batch normalization, and ReLU layers here

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=False):
        super().__init__()
        self.conv2d = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
        self.batchnorm2d = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.batchnorm2d(self.conv2d(x)))

model = ConvBlock(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
input_tensor = torch.randn(32, 3, 64, 64)

output = model(input_tensor)
assert output.shape == (32, 16, 64, 64), f'Expected output shape (32, 16, 64, 64), got {output.shape}'

assert isinstance(model.conv2d, nn.Conv2d), 'conv2d should be an instance of Conv2d'

assert isinstance(model.batchnorm2d, nn.BatchNorm2d), 'batchnorm2d should be an instance of BatchNorm2d'

assert isinstance(model.relu, nn.ReLU), 'relu should be an instance of ReLU'

class ConvBlock(Model):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=False):
        super().__init__()
        # Define the convolutional, batch normalization, and ReLU layers here

class ConvBlock(Model):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, bias=False):
        super().__init__()
        self.conv2d = layers.Conv2D(out_channels, kernel_size=kernel_size, strides=stride, padding='same', use_bias=bias, input_shape=(None, None, in_channels))
        self.batchnorm2d = layers.BatchNormalization()
        self.relu = layers.ReLU()

    def call(self, x):
        return self.relu(self.batchnorm2d(self.conv2d(x)))

model = ConvBlock(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding='same')
input_tensor = tf.random.normal([32, 64, 64, 3])

output = model(input_tensor)
assert output.shape == (32, 64, 64, 16), f'Expected output shape (32, 64, 64, 16), got {output.shape}'

assert isinstance(model.conv2d, layers.Conv2D), 'conv2d should be an instance of Conv2D'

assert isinstance(model.batchnorm2d, layers.BatchNormalization), 'batchnorm2d should be an instance of BatchNormalization'

assert isinstance(model.relu, layers.ReLU), 'relu should be an instance of ReLU'


class TwoLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super().__init__()
        # Define the layers here

class TwoLayerNet(nn.Module):
    def __init__(self, D_in, H, D_out):
        super().__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x):
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

model = TwoLayerNet(D_in=100, H=50, D_out=10)
input_tensor = torch.randn(64, 100)

output = model(input_tensor)
assert output.shape == (64, 10), f'Expected output shape (64, 10), got {output.shape}'

assert isinstance(model.linear1, nn.Linear), 'First layer should be a Linear layer'

assert isinstance(model.linear2, nn.Linear), 'Second layer should be a Linear layer'

import tensorflow as tf
from tensorflow.keras import layers, Model

class TwoLayerNet(Model):
    def __init__(self, D_in, H, D_out):
        super().__init__()
        # Define the layers here

class TwoLayerNet(Model):
    def __init__(self, D_in, H, D_out):
        super().__init__()
        self.linear1 = layers.Dense(H, input_shape=(D_in,), activation='relu')
        self.linear2 = layers.Dense(D_out)

    def call(self, x):
        x = self.linear1(x)
        y_pred = self.linear2(x)
        return y_pred

model = TwoLayerNet(D_in=100, H=50, D_out=10)
input_tensor = tf.random.normal([64, 100])

assert isinstance(model.linear1, layers.Dense), 'First layer should be a Dense layer'

assert isinstance(model.linear2, layers.Dense), 'Second layer should be a Dense layer'


torch.manual_seed(0)
square_matrix = torch.randn(3, 3)
# Calculate eigenvalues
# eigvals =

assert np.allclose(eigvals.numpy(), expected_values, atol=1e-4), 'Eigenvalue calculation does not match expected values'

import tensorflow as tf
import numpy as np

tf.random.set_seed(0)
square_matrix = tf.random.normal([3, 3], seed=0)
# Calculate eigenvalues
# eigvals =


conv1 = torch.nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)
# Initialize weights
# torch.nn.init.???(conv1.weight)

assert conv1.weight.std().item() > 0, 'Weights are not initialized'

assert abs(conv1.weight.mean().item()) < 0.01, 'Weight initialization is biased'

conv1 = tf.keras.layers.Conv2D(filters=16, kernel_size=3, input_shape=(None, None, 3))
# Initialize weights
# initializer = tf.keras.initializers.???
# conv1.kernel_initializer = initializer

initializer = tf.keras.initializers.GlorotUniform()
conv1.kernel_initializer = initializer

conv1.build(input_shape=(None, 32, 32, 3))

assert conv1.kernel.numpy().std() > 0, 'Weights are not initialized'

assert abs(conv1.kernel.numpy().mean()) < 0.01, 'Weight initialization is biased'


x = torch.tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
# Create a deep copy of x that does not require gradients
# y =

assert torch.allclose(x, y), 'Copied tensor values do not match original'

assert y.requires_grad == False, 'Copied tensor should not require gradients'

assert y.grad is None, 'Gradients were not detached from copied tensor'

x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
# Create a deep copy of x
# y =

assert tf.reduce_all(tf.equal(x, y)), 'Copied tensor values do not match original'

assert y.dtype == x.dtype, 'Copied tensor dtype does not match original'


x = torch.randn(2, 3, 4)
# Permute dimensions to the order (1, 2, 0)
# y =

assert y.shape == expected_shape, f'Output shape {y.shape} does not match expected {expected_shape}'

assert torch.allclose(y[0, 0], x[:, 0, 0]), 'Permuted tensor values do not match expected'

x = tf.random.normal(shape=(2, 3, 4))
# Permute dimensions to the order (1, 2, 0)
# y =

assert tf.reduce_all(tf.equal(y[0, 0], x[:, 0, 0])), 'Permuted tensor values do not match expected'


# Create a tensor with values drawn from a normal distribution with mean=2 and std=4
# x =

x = torch.normal(mean=2, std=4, size=(100,))

import numpy as np
check_mean = 2
check_std = 4
tol = 1.0 # allow some tolerance due to randomness

assert np.abs(x.mean().item() - check_mean) < tol, f'Sample mean {x.mean().item():.2f} deviates too much from {check_mean}'

assert np.abs(x.std().item() - check_std) < tol, f'Sample std {x.std().item():.2f} deviates too much from {check_std}'

x = tf.random.normal(shape=(100,), mean=2, stddev=4)

assert np.abs(tf.reduce_mean(x).numpy() - check_mean) < tol, f'Sample mean {tf.reduce_mean(x).numpy():.2f} deviates too much from {check_mean}'

assert np.abs(tf.math.reduce_std(x).numpy() - check_std) < tol, f'Sample std {tf.math.reduce_std(x).numpy():.2f} deviates too much from {check_std}'


n = 5
indices = torch.randint(0, n, size=(4, 7))
# Create one-hot vectors
# one_hot =

one_hot = torch.nn.functional.one_hot(indices, n)

assert one_hot.shape == (4, 7, n), f'Incorrect output shape: {one_hot.shape}, expected (4, 7, {n})'

assert torch.allclose(one_hot.sum(dim=-1), torch.tensor(1)), 'Sum of each one-hot vector should be 1'

assert torch.allclose(one_hot.argmax(dim=-1), indices), 'Incorrect one-hot encoding'

n = 5
indices = tf.random.uniform(shape=(4, 7), maxval=n, dtype=tf.int64)
# Create one-hot vectors
# one_hot =

assert tf.reduce_all(tf.equal(tf.reduce_sum(one_hot, axis=-1), 1)), 'Sum of each one-hot vector should be 1'

assert tf.reduce_all(tf.equal(tf.argmax(one_hot, axis=-1), indices)), 'Incorrect one-hot encoding'


lens = torch.LongTensor([3, 5, 4])
max_len = lens.max().item()
# Create a mask tensor
# mask =

mask = torch.arange(max_len).expand(len(lens), max_len) < lens.unsqueeze(1)
mask = mask.long()

assert isinstance(lens, torch.LongTensor), f'Input lens should be a torch.LongTensor, got {type(lens)}'

assert isinstance(mask, torch.LongTensor), f'Output mask should be a torch.LongTensor, got {type(mask)}'

assert mask.shape == expected_mask.shape, f'Incorrect mask shape: {mask.shape}, expected {expected_mask.shape}'

assert torch.all(mask == expected_mask), 'Mask values do not match expected'

lens = tf.constant([3, 5, 4], dtype=tf.int64)
max_len = tf.reduce_max(lens)
# Create a mask tensor
# mask =

mask = tf.sequence_mask(lens, maxlen=max_len, dtype=tf.int64)

assert lens.dtype == tf.int64, f'Input lens should have dtype int64, got {lens.dtype}'

assert mask.dtype == tf.int64, f'Output mask should have dtype int64, got {mask.dtype}'

assert tf.reduce_all(tf.equal(mask, expected_mask)), 'Mask values do not match expected'


t = torch.tensor([True, False, True, False])
# Convert to integer tensor
# t_integer =

expected_output = torch.tensor([1, 0, 1, 0])

assert isinstance(t, torch.BoolTensor), f'Input tensor should be of type torch.BoolTensor, got {t.dtype}'

assert isinstance(t_integer, torch.LongTensor), f'Output tensor should be of type torch.LongTensor, got {t_integer.dtype}'

assert torch.all(t_integer == expected_output), 'Converted values do not match expected'

t = tf.constant([True, False, True, False])
# Convert to integer tensor
# t_integer =

expected_output = tf.constant([1, 0, 1, 0], dtype=tf.int64)

assert t.dtype == tf.bool, f'Input tensor should have dtype bool, got {t.dtype}'

assert t_integer.dtype == tf.int64, f'Output tensor should have dtype int64, got {t_integer.dtype}'

assert tf.reduce_all(tf.equal(t_integer, expected_output)), 'Converted values do not match expected'


t = torch.randn(3, 4)
mean = 0.0
std = 0.1
# Add Gaussian noise to the tensor
# noisy_t =

noisy_t = t + torch.randn_like(t) * std + mean

check_mean = 0.0
check_std = 0.1
tol = 0.2 # allow some tolerance due to randomness

assert noisy_t.shape == t.shape, f'Shape of noisy tensor {noisy_t.shape} does not match original {t.shape}'

assert abs(noisy_t.mean().item() - (t.mean().item() + check_mean)) < tol, 'Mean of noisy tensor deviates too much from expected'

assert abs(noisy_t.std().item() - (t.std().item() + check_std)) < tol, 'Std of noisy tensor deviates too much from expected'

t = tf.random.normal(shape=(3, 4))
mean = 0.0
std = 0.1
# Add Gaussian noise to the tensor
# noisy_t =

noise = tf.random.normal(shape=t.shape, mean=mean, stddev=std)
noisy_t = t + noise

assert abs(tf.reduce_mean(noisy_t).numpy() - (tf.reduce_mean(t).numpy() + check_mean)) < tol, 'Mean of noisy tensor deviates too much from expected'

assert abs(tf.math.reduce_std(noisy_t).numpy() - (tf.math.reduce_std(t).numpy() + check_std)) < tol, 'Std of noisy tensor deviates too much from expected'


import torch
from torch.autograd import Variable

x = torch.tensor([[1, 2, 3], [4, 5, 6]])
# Create a 6 x 2 x 3 tensor by repeating x along a new dimension
# y =

assert y.shape == (6, 2, 3), f'Output shape {y.shape} does not match expected (6, 2, 3)'

assert torch.all(y == expected_output), 'Output values do not match expected'

x = tf.constant([[1, 2, 3], [4, 5, 6]])
# Create a 6 x 2 x 3 tensor by repeating x along a new dimension
# y =

y = tf.tile(tf.expand_dims(x, axis=0), multiples=[6, 1, 1])

assert tf.reduce_all(tf.equal(y, expected_output)), 'Output values do not match expected'


a = torch.tensor([1, 2, 3])
b = torch.tensor([1, 2])
c = torch.tensor([1])
# Create a padded tensor from the sequences
# padded_tensor =

padded_tensor = torch.nn.utils.rnn.pad_sequence([a, b, c], batch_first=True)

assert padded_tensor.shape == (3, 3), f'Output shape {padded_tensor.shape} does not match expected (3, 3)'

assert torch.equal(padded_tensor, expected_output), 'Output values do not match expected'

assert isinstance(padded_tensor, torch.Tensor), f'Output should be a torch.Tensor, got {type(padded_tensor)}'

a = tf.constant([1, 2, 3])
b = tf.constant([1, 2])
c = tf.constant([1])
# Create a padded tensor from the sequences
# padded_tensor =

padded_tensor = tf.keras.preprocessing.sequence.pad_sequences([a.numpy(), b.numpy(), c.numpy()], padding='post')
padded_tensor = tf.constant(padded_tensor)

assert tf.reduce_all(tf.equal(padded_tensor, expected_output)), 'Output values do not match expected'

assert isinstance(padded_tensor, tf.Tensor), f'Output should be a tf.Tensor, got {type(padded_tensor)}'


t = torch.tensor([[1, 2, 3],
                  [4, 5, 6],
                  [7, 8, 9]])
# Create a mask tensor where values are greater than 3 and less than 7
# mask =

expected_output = torch.tensor([[False, False, False],
                                [True,  True,  True],
                                [False, False, False]])

assert mask.shape == t.shape, f'Mask shape {mask.shape} does not match input tensor shape {t.shape}'

assert torch.equal(mask, expected_output), 'Mask values do not match expected'

assert mask.dtype == torch.bool, f'Mask should be of dtype torch.bool, got {mask.dtype}'

t = tf.constant([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]])
# Create a mask tensor where values are greater than 3 and less than 7
# mask =

expected_output = tf.constant([[False, False, False],
                               [True,  True,  True],
                               [False, False, False]])

assert tf.reduce_all(tf.equal(mask, expected_output)), 'Mask values do not match expected'

assert mask.dtype == tf.bool, f'Mask should be of dtype tf.bool, got {mask.dtype}'


import torch
from torch.utils.data import Dataset, DataLoader, RandomSampler

class MyDataset(Dataset):
    def __init__(self, N):
        self.N = N
        self.x = torch.rand(self.N, 10)
        self.y = torch.randint(0, 3, (self.N,))

    def __len__(self):
        return self.N

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]

# Create an instance of MyDataset with N samples
N = 100
ds = MyDataset(N)

# Create a DataLoader to sample M samples with replacement
M = 50
# sampler =
# dl =

sampler = RandomSampler(ds, replacement=True, num_samples=M)
dl = DataLoader(ds, batch_size=5, sampler=sampler)

num_unique_samples = len(set((x.numpy().tobytes(), y.numpy().tobytes()) for x, y in dl))

assert len(ds) == N, f'Dataset size should be {N}, got {len(ds)}'

assert num_unique_samples < M, 'Sampling with replacement should result in fewer unique samples than M'

assert all(x.shape == (5, 10) and y.shape == (5,) for x, y in dl), 'Batch shapes should be (4, 10) and (4,)'

# Create a dummy dataset
N = 100
x = tf.random.uniform((N, 10))
y = tf.random.uniform((N,), maxval=3, dtype=tf.int32)
ds = tf.data.Dataset.from_tensor_slices((x, y))

# Sample M samples with replacement
M = 50
# ds_sampled =

ds_sampled = ds.repeat().take(M)
ds_sampled = ds_sampled.batch(5)

num_batches = len(list(ds_sampled))
num_samples = sum(x.shape[0] for x, _ in ds_sampled)

assert num_samples >= M, f'Number of samples should be at least {M}, got {num_samples}'

assert all(x.shape[0] == 5 for x, _ in ds_sampled), 'Batch size should be 4'

assert num_batches == (M + 3) // 5, f'Number of batches should be {(M + 3) // 5}, got {num_batches}'


vocab = ['a', 'b', 'c', 'd']
data = [['a', 'c', 'd'], ['d', 'a', 'b']]
# Create a PyTorch Module of StringLookup
class StringLookup(torch.nn.Module):
    def __init__(self, vocabulary):
        super().__init__()
        # Initialize the lookup table

        # self.lookup_table =

    def forward(self, x):
        # Perform the lookup operation
        pass
        # return
# Create an instance of StringLookup

# layer =

# Apply the StringLookup layer to the data
# output =

class StringLookup(torch.nn.Module):
    def __init__(self, vocabulary):
        super().__init__()
        self.lookup_table = {word: idx for idx, word in enumerate(vocabulary)}

    def forward(self, x):
        indices = [[self.lookup_table[word] for word in seq] for seq in x]
        return torch.tensor(indices)

layer = StringLookup(vocabulary=vocab)
output = layer(data)
output

expected_output = torch.tensor([[0, 2, 3], [3, 0, 1]])

assert isinstance(layer, torch.nn.Module), 'StringLookup should be a torch.nn.Module'

assert torch.allclose(output, expected_output), 'Output does not match the expected result'

assert output.dtype == torch.long, f'Output should have dtype torch.long, got {output.dtype}'

vocab = ['a', 'b', 'c', 'd']
data = tf.constant([['a', 'c', 'd'], ['d', 'a', 'b']])

# Create a StringLookup layer
# layer =

# Apply the StringLookup layer to the data
# output =

layer = tf.keras.layers.StringLookup(vocabulary=vocab)
output = layer(data)

expected_output = tf.constant([[1, 3, 4], [4, 1, 2]],dtype=tf.int64)
print(output)

assert isinstance(layer, tf.keras.layers.StringLookup), 'Layer should be an instance of tf.keras.layers.StringLookup'

assert tf.reduce_all(tf.equal(output, expected_output)), 'Output does not match the expected result'

assert output.dtype == tf.int64, f'Output should have dtype int64, got {output.dtype}'


valid_cols = []
for col_idx in range(A.size(1)):
    if not torch.all(A[:, col_idx] == 0):
        valid_cols.append(col_idx)
A = A[:, valid_cols]

assert A.shape == expected_output.shape, f'Output shape {A.shape} does not match expected shape {expected_output.shape}'

assert torch.allclose(A, expected_output), 'Output values do not match the expected result'

assert all(torch.any(A[:, col_idx] != 0) for col_idx in range(A.size(1))), 'Output still contains columns filled with zeros'

valid_cols = tf.where(tf.reduce_any(A != 0, axis=0))
A = tf.gather(A, valid_cols, axis=1)
A= tf.squeeze(A, axis=-1)

assert tf.reduce_all(A == expected_output), 'Output values do not match the expected result'

assert tf.reduce_all(tf.reduce_any(A != 0, axis=0)), 'Output still contains columns filled with zeros'


n = 2
d = 5
rand_mat = torch.rand(n, d)
percent_ones = 0.25

# Create an integer tensor where percent_ones of elements are 1s and the rest are 0s
# k =
# k_th_quant =
# bool_tensor =
# desired_tensor =

k = round(percent_ones * d)
k_th_quant = torch.topk(rand_mat, k, largest=False)[0][:, -1:]
bool_tensor = rand_mat <= k_th_quant
desired_tensor = torch.where(bool_tensor, torch.tensor(1), torch.tensor(0))

num_ones = torch.sum(desired_tensor)
tol = 0.2 # Tolerance for percentage deviation

assert desired_tensor.dtype == torch.int64, f'Output should have dtype torch.int64, got {desired_tensor.dtype}'

assert desired_tensor.shape == (n, d), f'Output shape should be ({n}, {d}), got {desired_tensor.shape}'

assert desired_tensor.unique().tolist() == [0, 1], 'Output should contain only 0s and 1s'

assert abs(num_ones / (n * d) - percent_ones) < tol, f'Percentage of 1s should be close to {percent_ones}'

n = 2
d = 5
rand_mat = tf.random.uniform(shape=(n, d))
percent_ones = 0.25

# Create an integer tensor where percent_ones of elements are 1s and the rest are 0s
# k =
# k_th_quant =
# bool_tensor =
# desired_tensor =

k = round(percent_ones * d)
k_th_quant = tf.sort(rand_mat, direction='ASCENDING')[..., k - 1:k]
bool_tensor = rand_mat <= k_th_quant
desired_tensor = tf.cast(bool_tensor, tf.int32)

tf.experimental.numpy.experimental_enable_numpy_behavior()
num_ones = tf.reduce_sum(desired_tensor)
tol = 0.2 # Tolerance for percentage deviation

assert desired_tensor.dtype == tf.int32, f'Output should have dtype int32, got {desired_tensor.dtype}'

assert sorted(tf.unique(tf.squeeze(desired_tensor.reshape(n*d,-1))).y.numpy()) == [0, 1], 'Output should contain only 0s and 1s'


x = torch.arange(24).view(4, 3, 2)
ids = torch.randint(0, 3, size=(4, 1))

# Gather rows from x using indices in ids
# idx =
# result =

idx = ids.repeat(1, 2).view(4, 1, 2)
result = torch.gather(x, 1, idx)

expected_output = x[torch.arange(4), ids.squeeze()]

assert result.shape == (4, 1, 2), f'Output shape should be (4, 1, 2), got {result.shape}'

assert torch.allclose(result.squeeze(), expected_output), 'Output values do not match the expected result'

assert result.dtype == x.dtype, f'Output dtype {result.dtype} does not match input dtype {x.dtype}'

x = tf.reshape(tf.range(24), (4, 3, 2))
ids = tf.random.uniform(shape=(4, 1), maxval=3, dtype=tf.int32)

# Gather rows from x using indices in ids
# idx =
# result =

idx = tf.tile(ids, [1, 2])
result = tf.gather(x, idx, batch_dims=1)

expected_output = tf.gather(x, ids, batch_dims=1)

assert tf.squeeze(result).shape == (4, 2, 2), f'Output shape should be (4, 2, 2), got {result.shape}'

assert tf.reduce_all(result == expected_output), 'Output values do not match the expected result'


A = torch.randint(2, (10,))
B = torch.randint(2, (10,))

# Count the number of equal elements in A and B
# num_equal =

assert isinstance(num_equal, torch.Tensor), f'Output should be a torch.Tensor, got {type(num_equal)}'

assert num_equal.dtype == torch.int64, f'Output dtype should be torch.int64, got {num_equal.dtype}'

assert num_equal.item() == expected_output.item(), f'Expected {expected_output.item()} equal elements, got {num_equal.item()}'

A = tf.random.uniform(shape=(10,), maxval=2, dtype=tf.int32)
B = tf.random.uniform(shape=(10,), maxval=2, dtype=tf.int32)

# Count the number of equal elements in A and B
# num_equal =

num_equal = tf.reduce_sum(tf.cast(tf.equal(A, B), tf.int32))

expected_output = tf.reduce_sum(tf.cast(tf.equal(A, B), tf.int32))

assert isinstance(num_equal, tf.Tensor), f'Output should be a tf.Tensor, got {type(num_equal)}'

assert num_equal.dtype == tf.int32, f'Output dtype should be tf.int32, got {num_equal.dtype}'

assert num_equal.numpy() == expected_output.numpy(), f'Expected {expected_output.numpy()} equal elements, got {num_equal.numpy()}'


A = torch.zeros(2, dtype=torch.float32)  # sum vector
B = torch.tensor([0, 0, 1, 1])  # contribution vector
C = torch.tensor([20.0, 30.0, 40.0, 10.0])  # value vector

# Sum values in C according to indices in B and place results in A
# A.index_add_(...)

expected_output = torch.tensor([50.0, 50.0])

assert isinstance(A, torch.Tensor), f'A should be a torch.Tensor, got {type(A)}'

assert A.shape == (2,), f'A should have shape (2,), got {A.shape}'

assert torch.allclose(A, expected_output), f'Expected A to be {expected_output}, got {A}'

A = tf.zeros(2, dtype=tf.float32)  # sum vector
B = tf.constant([0, 0, 1, 1])  # contribution vector
C = tf.constant([20.0, 30.0, 40.0, 10.0])  # value vector

# Sum values in C according to indices in B and place results in A
# A =

A = tf.tensor_scatter_nd_add(A, tf.expand_dims(B, axis=-1), C)

expected_output = tf.constant([50.0, 50.0])

assert isinstance(A, tf.Tensor), f'A should be a tf.Tensor, got {type(A)}'

assert tf.reduce_all(tf.equal(A, expected_output)), f'Expected A to be {expected_output}, got {A}'


x = torch.tensor([
    [1, 2, 3, 4, 3, 3, 4],
    [1, 6, 3, 5, 3, 5, 4]
], dtype=torch.long)

# Zero out duplicate values in each row of x
# result =

y, indices = x.sort(dim=-1)
y[:, 1:] *= ((y[:, 1:] - y[:, :-1]) != 0).long()
indices = indices.sort(dim=-1)[1]
result = torch.gather(y, 1, indices)

assert isinstance(result, torch.Tensor), f'result should be a torch.Tensor, got {type(result)}'

assert result.shape == x.shape, f'result shape {result.shape} does not match input shape {x.shape}'

assert torch.all(result == expected_output), f'Expected result:\n{expected_output}\nGot:\n{result}'

x = tf.constant([
    [1, 2, 3, 4, 3, 3, 4],
    [1, 6, 3, 5, 3, 5, 4]
], dtype=tf.int64)

# Zero out duplicate values in each row of x
# result =

y = tf.sort(x, axis=-1)
indices = tf.argsort(x, axis=-1)
mask = tf.cast(tf.not_equal(y[:, 1:], y[:, :-1]), tf.int64)
y = tf.concat([y[:, :1], y[:, 1:] * mask], axis=-1)
result = tf.gather(y, indices, batch_dims=1)

assert isinstance(result, tf.Tensor), f'result should be a tf.Tensor, got {type(result)}'

assert tf.reduce_all(tf.equal(result, expected_output)), f'Expected result:\n{expected_output}\nGot:\n{result}'


old_values = torch.tensor([1, 2, 3, 4, 5, 5, 2, 3, 3, 2], dtype=torch.int32)
old_new_value = torch.tensor([[2, 22], [3, 33], [6, 66]], dtype=torch.int32)

# Replace old values with new values in old_values
# result =

old_new_dict = dict(old_new_value.tolist())
result = old_values.clone()
for old_val, new_val in old_new_dict.items():
    result[old_values == old_val] = new_val

expected_output = torch.tensor([1, 22, 33, 4, 5, 5, 22, 33, 33, 22], dtype=torch.int32)

assert result.shape == old_values.shape, f'result shape {result.shape} does not match input shape {old_values.shape}'

old_values = tf.constant([1, 2, 3, 4, 5, 5, 2, 3, 3, 2], dtype=tf.int32)
old_new_value = tf.constant([[2, 22], [3, 33], [6, 66]], dtype=tf.int32)

# Replace old values with new values in old_values
# result =

old_new_dict = dict(old_new_value.numpy())
result = tf.Variable(old_values)
for old_val, new_val in old_new_dict.items():
    result.assign(tf.where(old_values == old_val, new_val, result))

assert isinstance(result, tf.Variable), f'result should be a tf.Variable, got {type(result)}'


import torch
import torch.nn as nn
import numpy as np
import random

random_image = []
for x in range(1, 946):
    random_image.append(random.randint(0, 255))

random_image_arr = np.array(random_image)

# Create a Conv2d layer
# conv2 =

# Reshape the image to the required input shape
# image_d =

# Apply the convolution operation
# fc =

conv2 = nn.Conv2d(1, 18, kernel_size=3, stride=1, padding=1)
image_d = torch.FloatTensor(np.asarray(random_image_arr.reshape(1, 1, 27, 35)))
fc = conv2(image_d)

assert isinstance(conv2, nn.Conv2d), f'conv2 should be an instance of nn.Conv2d, got {type(conv2)}'

assert isinstance(image_d, torch.Tensor), f'image_d should be a torch.Tensor, got {type(image_d)}'

assert image_d.shape == (1, 1, 27, 35), f'Expected image_d shape (1, 1, 27, 35), got {image_d.shape}'

assert isinstance(fc, torch.Tensor), f'fc should be a torch.Tensor, got {type(fc)}'

assert fc.shape == expected_output_shape, f'Expected output shape {expected_output_shape}, got {fc.shape}'

import tensorflow as tf
import numpy as np
import random

random_image = []
for x in range(1, 946):
    random_image.append(random.randint(0, 255))

random_image_arr = np.array(random_image, dtype=np.float32)

# Create a Conv2D layer
# conv2 =

# Reshape the image to the required input shape
# image_d =

# Apply the convolution operation
# fc =

conv2 = tf.keras.layers.Conv2D(18, kernel_size=3, strides=1, padding='same')
image_d = tf.reshape(random_image_arr, [1, 27, 35, 1])
fc = conv2(image_d)

assert isinstance(conv2, tf.keras.layers.Conv2D), f'conv2 should be an instance of tf.keras.layers.Conv2D, got {type(conv2)}'

assert isinstance(image_d, tf.Tensor), f'image_d should be a tf.Tensor, got {type(image_d)}'

assert image_d.shape == (1, 27, 35, 1), f'Expected image_d shape (1, 27, 35, 1), got {image_d.shape}'

assert isinstance(fc, tf.Tensor), f'fc should be a tf.Tensor, got {type(fc)}'


t = torch.rand(2, 3)

# Calculate the 2-norm along dimension 1
# norm_2 =

norm_2 = t.norm(dim=1, p=2)
# Alternative solutions:
# norm_2 = torch.linalg.norm(t, dim=1, ord=2)
# norm_2 = t.pow(2).sum(dim=1).sqrt()

expected_output = torch.sqrt(torch.sum(t**2, dim=1))

assert isinstance(norm_2, torch.Tensor), f'norm_2 should be a torch.Tensor, got {type(norm_2)}'

assert norm_2.shape == (2,), f'Expected norm_2 shape: (2,), got: {norm_2.shape}'

assert torch.allclose(norm_2, expected_output), f'Expected output: {expected_output}, got: {norm_2}'

t = tf.random.uniform(shape=(2, 3))

# Calculate the 2-norm along dimension 1
# norm_2 =

norm_2 = tf.norm(t, ord=2, axis=1)
# Alternative solution:
# norm_2 = tf.sqrt(tf.reduce_sum(tf.square(t), axis=1))

expected_output = tf.sqrt(tf.reduce_sum(tf.square(t), axis=1))

assert isinstance(norm_2, tf.Tensor), f'norm_2 should be a tf.Tensor, got {type(norm_2)}'

assert tf.reduce_all(tf.abs(norm_2 - expected_output) < 1e-6), f'Expected output: {expected_output}, got: {norm_2}'


A = torch.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
B = torch.tensor([2, 5, 6, 8, 12, 15, 16])

# Create a comparison matrix
# comparison_matrix =

# Count the number of elements from B that exist in A
# exists =

# Count the number of elements from B that do not exist in A
# not_exist =

comparison_matrix = B.unsqueeze(1) == A.unsqueeze(0)
matches_per_element = comparison_matrix.sum(dim=1)

exists = (matches_per_element > 0).sum().item()
not_exist = len(B) - exists

assert isinstance(exists, int), f'exists should be an int, got {type(exists)}'

assert isinstance(not_exist, int), f'not_exist should be an int, got {type(not_exist)}'

assert exists == expected_exists, f'Expected exists: {expected_exists}, got: {exists}'

assert not_exist == expected_not_exist, f'Expected not_exist: {expected_not_exist}, got: {not_exist}'

A = tf.constant([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
B = tf.constant([2, 5, 6, 8, 12, 15, 16])

# Create a comparison matrix
# comparison_matrix =

# Count the number of elements from B that exist in A
# exists =

# Count the number of elements from B that do not exist in A
# not_exist =

comparison_matrix = tf.equal(tf.expand_dims(B, 1), tf.expand_dims(A, 0))
matches_per_element = tf.reduce_sum(tf.cast(comparison_matrix, tf.int32), axis=1)

exists = tf.reduce_sum(tf.cast(matches_per_element > 0, tf.int32)).numpy()
not_exist = len(B) - exists

import numpy as np
expected_exists = 4
expected_not_exist = 3

assert isinstance(exists, np.int32), f'exists should be an int, got {type(exists)}'

assert isinstance(not_exist, np.int64), f'not_exist should be an int, got {type(not_exist)}'


# Create a 3-dimensional tensor of shape (a, b, c)
tensor = torch.rand(4, 5, 6)

# Create a list of indices of length a, each in the range [0, b)
B = [1, 3, 2, 4]

# Slice the tensor along the first axis using the indices
# output =

output = tensor[torch.arange(len(B)), B]

expected_output = torch.stack([tensor[i, B[i]] for i in range(len(B))])

assert isinstance(output, torch.Tensor), f'output should be a torch.Tensor, got {type(output)}'

assert output.shape == (len(B), tensor.shape[-1]), f'Expected output shape: ({len(B)}, {tensor.shape[-1]}), got: {output.shape}'

assert torch.allclose(output, expected_output), f'Expected output:\n{expected_output}\nGot:\n{output}'

# Create a 3-dimensional tensor of shape (a, b, c)
tensor = tf.random.uniform(shape=(4, 5, 6))

# Create a list of indices of length a, each in the range [0, b)
B = [1, 3, 2, 4]

# Slice the tensor along the first axis using the indices
# output =

output = tf.gather_nd(tensor, tf.stack([tf.range(len(B)), B], axis=1))

expected_output = tf.stack([tensor[i, B[i]] for i in range(len(B))])

assert isinstance(output, tf.Tensor), f'output should be a tf.Tensor, got {type(output)}'

assert tf.reduce_all(tf.abs(output - expected_output) < 1e-6), f'Expected output:\n{expected_output}\nGot:\n{output}'


v1 = torch.tensor([2, 3, 3])
v2 = torch.tensor([1, 2, 6, 2, 3, 10, 4, 6, 4])

# Mask to find elements in v2 that are not in v1
mask = ~torch.isin(v2, v1)
v2_without_v1 = v2[mask]

# Get unique elements and their first indices
# unique_indices =

def get_unique_elements_first_idx(tensor):
    # sort tensor
    sorted_tensor, indices = torch.sort(tensor)
    # find position of jumps
    unique_mask = torch.cat((torch.tensor([True]), sorted_tensor[1:] != sorted_tensor[:-1]))
    return indices[unique_mask]

unique_indices = get_unique_elements_first_idx(v2_without_v1)

expected_unique_indices = torch.tensor([0, 3, 1, 2])
expected_unique_values = v2[mask][expected_unique_indices]

assert isinstance(unique_indices, torch.Tensor), f'unique_indices should be a torch.Tensor, got {type(unique_indices)}'

assert torch.all(unique_indices == expected_unique_indices), f'Expected unique_indices: {expected_unique_indices}, got: {unique_indices}'

assert torch.all(v2[mask][unique_indices] == expected_unique_values), f'Expected unique values: {expected_unique_values}, got: {v2[mask][unique_indices]}'

v1 = tf.constant([2, 3, 3])
v2 = tf.constant([1, 2, 6, 2, 3, 10, 4, 6, 4])

# Mask to find elements in v2 that are not in v1
mask = ~tf.reduce_any(tf.equal(tf.expand_dims(v2, 1), tf.expand_dims(v1, 0)), axis=1)
v2_without_v1 = tf.boolean_mask(v2, mask)

# Get unique elements and their first indices
# unique_indices =

import tensorflow as tf

def get_unique_elements_first_idx(tensor):
    # sort tensor
    sorted_tensor = tf.sort(tensor)
    # indices of sorted tensor
    indices = tf.argsort(tensor)
    # find position of jumps
    unique_mask = tf.concat(
        ([True], sorted_tensor[1:] != sorted_tensor[:-1]), axis=0)
    return tf.boolean_mask(indices, unique_mask)

unique_indices = get_unique_elements_first_idx(v2_without_v1)

expected_unique_indices = tf.constant([0, 3, 1, 2])
expected_unique_values = tf.gather(v2_without_v1, expected_unique_indices)

assert isinstance(unique_indices, tf.Tensor), f'unique_indices should be a tf.Tensor, got {type(unique_indices)}'

assert tf.reduce_all(unique_indices == expected_unique_indices), f'Expected unique_indices: {expected_unique_indices}, got: {unique_indices}'

assert tf.reduce_all(tf.gather(v2_without_v1, unique_indices) == expected_unique_values), f'Expected unique values: {expected_unique_values}, got: {tf.gather(v2_without_v1, unique_indices)}'


class ValidatedArray:
    def __init__(self, array: torch.Tensor):
        self._array = array
        self._validate_array()

    def _validate_array(self):
        # Add validation logic here
        pass

    @property
    def array(self):
        self._validate_array()
        return self._array

class ValidatedArray:
    def __init__(self, array: torch.Tensor):
        self._array = array
        self._validate_array()

    def _validate_array(self):
        assert torch.allclose(self._array.sum(dim=-1), torch.ones(self._array.shape[:-1])), f'The last dim represents a categorical distribution. It must sum to one.'

    @property
    def array(self):
        self._validate_array()
        return self._array

assert isinstance(ValidatedArray(valid_array).array, torch.Tensor), 'ValidatedArray.array should return a torch.Tensor'

assert torch.allclose(ValidatedArray(valid_array).array, valid_array), 'ValidatedArray.array should return the original tensor if it is valid'

try:
    ValidatedArray(invalid_array)
    assert False, 'ValidatedArray should raise an AssertionError for an invalid tensor'
except AssertionError:
    pass

class ValidatedArray:
    def __init__(self, array: tf.Tensor):
        self._array = array
        self._validate_array()

    def _validate_array(self):
        # Add validation logic here
        pass

    @property
    def array(self):
        self._validate_array()
        return self._array

class ValidatedArray:
    def __init__(self, array: tf.Tensor):
        self._array = array
        self._validate_array()

    def _validate_array(self):
        assert tf.reduce_all(tf.abs(tf.reduce_sum(self._array, axis=-1) - 1.0) < 1e-6), f'The last dim represents a categorical distribution. It must sum to one.'

    @property
    def array(self):
        self._validate_array()
        return self._array

assert isinstance(ValidatedArray(valid_array).array, tf.Tensor), 'ValidatedArray.array should return a tf.Tensor'

assert tf.reduce_all(ValidatedArray(valid_array).array == valid_array), 'ValidatedArray.array should return the original tensor if it is valid'


# Create a 5D tensor
tensor = torch.rand(1, 3, 10, 40, 1)

# Specify the dimension and size for splitting
dimension = 3
size=10
# Split the tensor into 10 tensors
# split_tensors =

split_tensors = torch.split(tensor, 10, dim=3)
# convert the tuple of tensors to a tensor
split_tensors = torch.stack(split_tensors)

assert isinstance(split_tensors, torch.Tensor), 'split_tensors should be a torch.Tensor'

assert split_tensors.shape == expected_shape, f'Expected shape {expected_shape}, got {split_tensors.shape}'

assert torch.allclose(split_tensors[0], tensor[:, :, :, 0:10,:]), 'First split tensor does not match expected values'

assert torch.allclose(split_tensors[-1], tensor[:, :, :, 30:40,:]), 'Last split tensor does not match expected values'

# Create a 5D tensor
tensor = tf.random.uniform(shape=(1, 3, 10, 40, 1))

# Specify the dimension and size for splitting
dimension = 3
size = 10

# Split the tensor into smaller tensors
# split_tensors =

# Create a 5D tensor
tensor = tf.random.uniform(shape=(1, 3, 10, 40, 1))

# Specify the dimension and size for splitting
dimension = 3
size = 10

# Split the tensor into smaller tensors
split_tensors = tf.split(tensor, num_or_size_splits=tensor.shape[dimension]//size, axis=dimension)
split_tensors = tf.stack(split_tensors)

assert isinstance(split_tensors, tf.Tensor), 'split_tensors should be a tf.Tensor'

assert tf.reduce_all(split_tensors[0] == tensor[:, :, :, :10]), 'First split tensor does not match expected values'

assert tf.reduce_all(split_tensors[-1] == tensor[:, :, :, 30:40]), 'Last split tensor does not match expected values'


import torch
import torch.nn.functional as F

# Assume we have a batch of images and masks with different sizes
image_batch = [torch.rand(3, 32, 32), torch.rand(3, 64, 64), torch.rand(3, 48, 48)]
mask_batch = [torch.rand(32, 32), torch.rand(64, 64), torch.rand(48, 48)]

# Determine maximum height and width
# max_height =
# max_width =

# Pad the image batch
# image_batch =

# Pad the mask batch
# mask_batch =

# Determine maximum height and width
max_height = max([img.size(1) for img in image_batch])
max_width = max([img.size(2) for img in image_batch])

# Pad the image batch
image_batch = [F.pad(img, [0, max_width - img.size(2), 0, max_height - img.size(1)]) for img in image_batch]

# Pad the mask batch
mask_batch = [F.pad(mask, [0, max_width - mask.size(1), 0, max_height - mask.size(0)]) for mask in mask_batch]

assert isinstance(image_batch, list) and isinstance(mask_batch, list), 'image_batch and mask_batch should be lists'

assert all(isinstance(img, torch.Tensor) for img in image_batch), 'image_batch should contain torch.Tensor objects'

assert all(isinstance(mask, torch.Tensor) for mask in mask_batch), 'mask_batch should contain torch.Tensor objects'

assert len(image_batch) == len(mask_batch), 'image_batch and mask_batch should have the same length'

assert all(img.size() == (3, max_height, max_width) for img in image_batch), 'All images should have the same padded size'

assert all(mask.size() == (max_height, max_width) for mask in mask_batch), 'All masks should have the same padded size'

# Assume we have a batch of images and masks with different sizes
image_batch = [tf.random.uniform((32, 32, 3)), tf.random.uniform((64, 64, 3)), tf.random.uniform((48, 48, 3))]
mask_batch = [tf.random.uniform((32, 32)), tf.random.uniform((64, 64)), tf.random.uniform((48, 48))]

# Determine maximum height and width
# max_height =
# max_width =

# Pad the image batch
# image_batch =

# Pad the mask batch
# mask_batch =

# Determine maximum height and width
max_height = max([img.shape[0] for img in image_batch])
max_width = max([img.shape[1] for img in image_batch])

# Pad the image batch
image_batch = [tf.pad(img, [[0, max_height - img.shape[0]], [0, max_width - img.shape[1]], [0, 0]]) for img in image_batch]

# Pad the mask batch
mask_batch = [tf.pad(mask, [[0, max_height - mask.shape[0]], [0, max_width - mask.shape[1]]]) for mask in mask_batch]

assert all(isinstance(img, tf.Tensor) for img in image_batch), 'image_batch should contain tf.Tensor objects'

assert all(isinstance(mask, tf.Tensor) for mask in mask_batch), 'mask_batch should contain tf.Tensor objects'

assert all(img.shape == (max_height, max_width, 3) for img in image_batch), 'All images should have the same padded size'

assert all(mask.shape == (max_height, max_width) for mask in mask_batch), 'All masks should have the same padded size'


a = torch.tensor([1, 234, 54, 6543, 55, 776])
b = torch.tensor([234, 54])
c = torch.tensor([55, 776])

# Check if values in a are contained in b or c and create a boolean mask
# a_masked =

a_masked = sum(a == i for i in b).bool() + sum(a == i for i in c).bool()

expected_output = torch.tensor([False,  True,  True, False, True, True])

assert isinstance(a_masked, torch.Tensor), 'a_masked should be a torch.Tensor'

assert a_masked.dtype == torch.bool, 'a_masked should have dtype torch.bool'

assert a_masked.shape == a.shape, f'a_masked should have the same shape as a. Expected {a.shape}, got {a_masked.shape}'

assert torch.all(a_masked == expected_output), f'\nExpected:\n{expected_output}\nGot:\n{a_masked}'

a = tf.constant([1, 234, 54, 6543, 55, 776])
b = tf.constant([234, 54])
c = tf.constant([55, 776])

# Check if values in a are contained in b or c and create a boolean mask
# a_masked =

a_masked = tf.reduce_any(tf.equal(a[:, tf.newaxis], b), axis=1) | tf.reduce_any(tf.equal(a[:, tf.newaxis], c), axis=1)

expected_output = tf.constant([False,  True,  True, False, True, True])

assert isinstance(a_masked, tf.Tensor), 'a_masked should be a tf.Tensor'

assert a_masked.dtype == tf.bool, 'a_masked should have dtype tf.bool'

assert tf.reduce_all(tf.equal(a_masked, expected_output)), f'\nExpected:\n{expected_output}\nGot:\n{a_masked}'


x = torch.rand(20, 30, 40, 50)
y = torch.rand(20, 30, 40, 50)

# Merge x and y based on their channel norms
# B, C, H, W = x.size()
# z = torch.zeros_like(x)

# x_norm =
# y_norm =

# condition =

# Assign values to z based on the condition
# z[condition] =
# z[~condition] =

B, C, H, W = x.size()
z = torch.zeros_like(x)

x_norm = torch.norm(x, dim=(2, 3))
y_norm = torch.norm(y, dim=(2, 3))

condition = x_norm >= y_norm

z[condition] = x[condition]
z[~condition] = y[~condition]

assert isinstance(z, torch.Tensor), 'z should be a torch.Tensor'

assert z.shape == x.shape, f'z should have the same shape as x and y. Expected {x.shape}, got {z.shape}'

assert torch.allclose(z[x_norm >= y_norm], x[x_norm >= y_norm]), 'z should contain values from x where x_norm >= y_norm'

assert torch.allclose(z[x_norm < y_norm], y[x_norm < y_norm]), 'z should contain values from y where x_norm < y_norm'

x = tf.random.uniform((20, 30, 40, 50))
y = tf.random.uniform((20, 30, 40, 50))

# Merge x and y based on their channel norms
# B, C, H, W = x.shape
# z = tf.zeros_like(x)

# x_norm =
# y_norm =

# condition =

# Assign values to z based on the condition
# z = tf.where(condition, x, y)

B, C, H, W = x.shape
z = tf.zeros_like(x)

x_norm = tf.norm(x, axis=(2, 3))
y_norm = tf.norm(y, axis=(2, 3))

condition = x_norm >= y_norm

z = tf.where(condition[:, :, tf.newaxis, tf.newaxis], x, y)

assert isinstance(z, tf.Tensor), 'z should be a tf.Tensor'

assert tf.reduce_all(z[condition] == x[condition]), 'z should contain values from x where condition is True'

assert tf.reduce_all(z[~condition] == y[~condition]), 'z should contain values from y where condition is False'


a = torch.FloatTensor([[[1, 1, 1], [2, 2, 2]], [[9, 9, 9], [5, 5, 5]]])
b = torch.IntTensor([1, 0])

# Select rows from a based on indices in b
# out =

expected_output = torch.FloatTensor([[2, 2, 2], [9, 9, 9]])

assert isinstance(out, torch.Tensor), 'out should be a torch.Tensor'

assert out.shape == (a.size(0), a.size(2)), f'out should have shape ({a.size(0)}, {a.size(2)}), got {out.shape}'

assert torch.allclose(out, expected_output), f'\nExpected:\n{expected_output}\nGot:\n{out}'

a = tf.constant([[[1, 1, 1], [2, 2, 2]], [[9, 9, 9], [5, 5, 5]]], dtype=tf.float32)
b = tf.constant([1, 0], dtype=tf.int32)

# Select rows from a based on indices in b
# out =

out = tf.gather_nd(a, tf.stack([tf.range(tf.shape(a)[0]), b], axis=1))

expected_output = tf.constant([[2, 2, 2], [9, 9, 9]], dtype=tf.float32)

assert isinstance(out, tf.Tensor), 'out should be a tf.Tensor'

assert out.shape == (a.shape[0], a.shape[2]), f'out should have shape ({a.shape[0]}, {a.shape[2]}), got {out.shape}'

assert tf.reduce_all(tf.equal(out, expected_output)), f'\nExpected:\n{expected_output}\nGot:\n{out}'


indexes = torch.tensor([[0, 2, 1, 3],
                        [1, 0, 3, 2]])
t1 = torch.zeros_like(indexes).float()
t2 = torch.tensor([[0.1, 0.2, 0.3, 0.4],
                   [0.5, 0.6, 0.7, 0.8]])

# Re-assign values in t1 based on indexes and t2
# t1[...] = ...

assert isinstance(t1, torch.Tensor), 't1 should be a torch.Tensor'

assert t1.shape == indexes.shape == t2.shape, 't1, indexes, and t2 should have the same shape'

assert torch.allclose(t1, expected_output), f'Expected:\n{expected_output}\nGot:\n{t1}'

indexes = tf.constant([[0, 2, 1, 3],
                       [1, 0, 3, 2]])
t1 = tf.zeros_like(indexes, dtype=tf.float32)
t2 = tf.constant([[0.1, 0.2, 0.3, 0.4],
                  [0.5, 0.6, 0.7, 0.8]])

# Re-assign values in t1 based on indexes and t2
# t1 = ...

t1 = tf.gather(t2, indexes, batch_dims=1)

assert isinstance(t1, tf.Tensor), 't1 should be a tf.Tensor'

assert tf.reduce_all(tf.equal(t1, expected_output)), f'Expected:\n{expected_output}\nGot:\n{t1}'


tensor1 = torch.rand((4, 2, 3, 100))
tensor2 = torch.rand((4, 2, 3, 100))

# Calculate the Euclidean distance between corresponding vectors
# dist =

dist = (tensor1 - tensor2).pow(2).sum(3).sqrt()

assert isinstance(dist, torch.Tensor), 'dist should be a torch.Tensor'

assert dist.shape == expected_shape, f'dist should have shape {expected_shape}, got {dist.shape}'

assert torch.allclose(dist, torch.norm(tensor1 - tensor2, dim=3), atol=tol), 'dist values do not match expected values'

tensor1 = tf.random.uniform((4, 2, 3, 100))
tensor2 = tf.random.uniform((4, 2, 3, 100))

# Calculate the Euclidean distance between corresponding vectors
# dist =

dist = tf.sqrt(tf.reduce_sum(tf.square(tensor1 - tensor2), axis=3))

assert isinstance(dist, tf.Tensor), 'dist should be a tf.Tensor'

assert tf.reduce_all(tf.abs(dist - tf.norm(tensor1 - tensor2, axis=3)) < tol), 'dist values do not match expected values'


assert isinstance(v, torch.Tensor), 'v should be a torch.Tensor'

assert v.shape == (3, 3), 'v should have shape (3, 3)'

assert torch.allclose(v, expected_output), f'Expected:\n{expected_output}\nGot:\n{v}'

# Create the target tensor
v = tf.constant([[1., 2., 3.],
                 [4., 5., 6.],
                 [7., 8., 9.]])

# Create the source tensor
w = tf.constant([[10., 20., 30.],
                 [40., 50., 60.],
                 [70., 80., 90.]])

# Create the index tensor
index = tf.constant([0, 1, 0])

# Add rows from w to v based on index
# v =

v = tf.tensor_scatter_nd_add(v, tf.expand_dims(index, axis=1), w)

assert isinstance(v, tf.Tensor), 'v should be a tf.Tensor'

assert tf.reduce_all(tf.equal(v, expected_output)), f'Expected:\n{expected_output}\nGot:\n{v}'


# Create a sample tensor
a = torch.tensor([[4, 2, 5],
                  [1, 3, 2],
                  [9, 7, 8]])

# Sort the rows based on the last column
# ind =
# sorted_a =

ind = a[:, -1].argsort(dim=0)
sorted_a = a[ind]

assert isinstance(sorted_a, torch.Tensor), 'sorted_a should be a torch.Tensor'

assert sorted_a.shape == a.shape, f'sorted_a should have the same shape as a. Expected {a.shape}, got {sorted_a.shape}'

assert torch.allclose(sorted_a, expected_output), f'Expected:\n{expected_output}\nGot:\n{sorted_a}'

# Create a sample tensor
a = tf.constant([[4, 2, 5],
                 [1, 3, 2],
                 [9, 7, 8]])

# Sort the rows based on the last column
# ind =
# sorted_a =

ind = tf.argsort(a[:, -1])
sorted_a = tf.gather(a, ind)

assert isinstance(sorted_a, tf.Tensor), 'sorted_a should be a tf.Tensor'

assert tf.reduce_all(tf.equal(sorted_a, expected_output)), f'Expected:\n{expected_output}\nGot:\n{sorted_a}'


class LR(torch.nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.lr = torch.nn.Sequential(
            # Add your layers here
        )

    def forward(self, x):
        return self.lr(x)

class LR(torch.nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.lr = torch.nn.Sequential(
            torch.nn.Linear(n_features, 64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, 36, bias=False),
            torch.nn.Softmax(dim=1),
        )

    def forward(self, x):
        return self.lr(x)

n_features = 384
model = LR(n_features)
sample_input = torch.randn(1, n_features)

assert isinstance(model, torch.nn.Module), 'model should be an instance of torch.nn.Module'

assert len(list(model.parameters())) == 3, 'model should have 3 sets of parameters'

assert output.shape == (1, 36), f'Expected output shape (1, 36), got {output.shape}'

model = tf.keras.models.Sequential([
    # Add your layers here
])

model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(64, input_shape=(384,), activation="relu"),
    tf.keras.layers.Dense(36, activation="softmax", use_bias=False)
])

n_features = 384
sample_input = tf.random.normal((1, n_features))

assert isinstance(model, tf.keras.Sequential), 'model should be an instance of tf.keras.Sequential'

assert len(model.layers) == 2, 'model should have 2 layers'


import torch
import torch.nn as nn
import torch.nn.functional as F

class FirstM(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(20, 2)

    def forward(self, x):
        x = self.fc1(x)
        return x

class SecondM(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(20, 2)

    def forward(self, x):
        x = self.fc1(x)
        return x

class Combined_model(nn.Module):
    def __init__(self, modelA, modelB):
        super(Combined_model, self).__init__()
        self.modelA = modelA
        self.modelB = modelB

        # Add a classifier layer to combine the outputs
        # self.classifier = ...

    def forward(self, x1, x2):
        # Pass inputs through modelA and modelB
        # x1 = ...
        # x2 = ...

        # Concatenate the outputs
        # x = ...

        # Pass the concatenated output through the classifier
        # x = ...

        return x

class Combined_model(nn.Module):
    def __init__(self, modelA, modelB):
        super().__init__()
        self.modelA = modelA
        self.modelB = modelB
        self.classifier = nn.Linear(4, 2)

    def forward(self, x1, x2):
        x1 = self.modelA(x1)
        x2 = self.modelB(x2)
        x = torch.cat((x1, x2), dim=1)
        x = self.classifier(F.relu(x))
        return x

modelA = FirstM()
modelB = SecondM()
combined_model = Combined_model(modelA, modelB)

x1 = torch.randn(1, 20)
x2 = torch.randn(1, 20)

assert isinstance(combined_model, nn.Module), 'combined_model should be an instance of nn.Module'

assert len(list(combined_model.parameters())) == 6, 'combined_model should have 6 learnable parameters'

assert output.shape == (1, 2), f'Expected output shape (1, 2), but got {output.shape}'

class FirstM(tf.keras.Model):
    def __init__(self):
        super(FirstM, self).__init__()
        self.fc1 = tf.keras.layers.Dense(2, input_shape=(20,))

    def call(self, x):
        x = self.fc1(x)
        return x

class SecondM(tf.keras.Model):
    def __init__(self):
        super(SecondM, self).__init__()
        self.fc1 = tf.keras.layers.Dense(2, input_shape=(20,))

    def call(self, x):
        x = self.fc1(x)
        return x

class Combined_model(tf.keras.Model):
    def __init__(self, modelA, modelB):
        super(Combined_model, self).__init__()
        self.modelA = modelA
        self.modelB = modelB

        # Add a classifier layer to combine the outputs
        # self.classifier = ...

    def call(self, x1, x2):
        # Pass inputs through modelA and modelB
        # x1 = ...
        # x2 = ...

        # Concatenate the outputs
        # x = ...

        # Pass the concatenated output through the classifier
        # x = ...

        return x

class Combined_model(tf.keras.Model):
    def __init__(self, modelA, modelB):
        super(Combined_model, self).__init__()
        self.modelA = modelA
        self.modelB = modelB
        self.classifier = tf.keras.layers.Dense(2, activation='relu')

    def call(self, x1, x2):
        x1 = self.modelA(x1)
        x2 = self.modelB(x2)
        x = tf.concat((x1, x2), axis=1)
        x = self.classifier(x)
        return x

modelA = FirstM()
modelB = SecondM()
combined_model = Combined_model(modelA, modelB)

x1 = tf.random.normal((1, 20))
x2 = tf.random.normal((1, 20))
combined_model(x1, x2)

assert isinstance(combined_model, tf.keras.Model), 'combined_model should be an instance of tf.keras.Model'

assert len(combined_model.trainable_variables) == 6, 'combined_model should have 6 trainable variables'


multiplier = torch.from_numpy(np.array([1.0, -2.0, 5.0]))
source = torch.from_numpy(
    np.array([[0.5, -0.6],
              [-3.0, -2.0],
              [-4.0, 2.3]]))
target = torch.from_numpy(np.array([0, 1, 0]))

# Create an instance of CrossEntropyLoss with reduction='none'
# loss_fn =

# Compute the batch-specific loss
# batch_loss =

# Weigh the batch loss using the multiplier
# weighted_loss =

# Compute the mean weighted loss
# loss =

loss_fn = torch.nn.CrossEntropyLoss(reduction='none')
batch_loss = loss_fn(source, target)
weighted_loss = multiplier * batch_loss
loss = torch.sum(weighted_loss) / torch.sum(multiplier)

assert isinstance(loss_fn, torch.nn.CrossEntropyLoss), 'loss_fn should be an instance of torch.nn.CrossEntropyLoss'

assert loss_fn.reduction == 'none', 'loss_fn should have reduction set to \'none\''

assert torch.is_tensor(batch_loss), 'batch_loss should be a torch.Tensor'

assert batch_loss.shape == multiplier.shape, f'batch_loss should have the same shape as multiplier, got {batch_loss.shape} and {multiplier.shape}'

assert torch.isclose(loss, torch.tensor(expected_output,dtype=torch.float64), atol=1e-4), f'Expected loss {expected_output}, but got {loss}'

multiplier = tf.constant([1.0, -2.0, 5.0])
source = tf.constant([[0.5, -0.6],
                      [-3.0, -2.0],
                      [-4.0, 2.3]])
target = tf.constant([0, 1, 0])

# Create an instance of SparseCategoricalCrossentropy with reduction='none'
# loss_fn =

# Compute the batch-specific loss
# batch_loss =

# Weigh the batch loss using the multiplier
# weighted_loss =

# Compute the mean weighted loss
# loss =

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(reduction='none')
batch_loss = loss_fn(target, source)
weighted_loss = multiplier * batch_loss
loss = tf.reduce_sum(weighted_loss) / tf.reduce_sum(multiplier)

assert isinstance(loss_fn, tf.keras.losses.Loss), 'loss_fn should be an instance of tf.keras.losses.Loss'

assert tf.is_tensor(batch_loss), 'batch_loss should be a tf.Tensor'

assert tf.math.abs(loss - expected_output) < 1e-4, f'Expected loss {expected_output}, but got {loss}'


# Create a sample tensor A with shape [M, N]
M, N = 3, 4
A = torch.randn(M, N)

# Specify the number of repetitions K
K = 5

# Repeat A along a new dimension to create B with shape [M, K, N]
# B =

assert isinstance(B, torch.Tensor), 'B should be a torch.Tensor'

assert B.shape == expected_shape, f'B should have shape {expected_shape}, but got {B.shape}'

assert torch.allclose(B[:, 0, :], A), 'B[:, 0, :] should be equal to A'

assert torch.allclose(B[:, -1, :], A), 'B[:, -1, :] should be equal to A'

# Create a sample tensor A with shape [M, N]
M, N = 3, 4
A = tf.random.normal((M, N))

# Specify the number of repetitions K
K = 5

# Repeat A along a new dimension to create B with shape [M, K, N]
# B =

B = tf.repeat(tf.expand_dims(A, 1), K, axis=1)

assert isinstance(B, tf.Tensor), 'B should be a tf.Tensor'

assert tf.reduce_all(B[:, 0, :] == A), 'B[:, 0, :] should be equal to A'

assert tf.reduce_all(B[:, -1, :] == A), 'B[:, -1, :] should be equal to A'


def sum_by_index(labels, source):
    unique_labels, labels_inds, labels_counts = labels.unique(return_inverse=True, return_counts=True)
    out = torch.zeros((unique_labels.size(0), source.size(2)), dtype=source.dtype)

    # Add the embeddings from source to out based on labels_inds
    # out.index_add_(...)

    return out

labels = torch.tensor([[0, 1], [1, 2]])
source = torch.tensor(
          [[[0, 1],
           [1, 2]],

           [[2, 3],
            [3, 4]]],
        dtype=torch.float,
        requires_grad=True,
)

# Apply sum_by_index to get the result
# result =

def sum_by_index(labels, source):
    unique_labels, labels_inds, labels_counts = labels.unique(return_inverse=True, return_counts=True)
    out = torch.zeros((unique_labels.size(0), source.size(2)), dtype=source.dtype)
    out.index_add_(0, labels_inds.flatten(), source.flatten(0, 1))
    return out

labels = torch.tensor([[0, 1], [1, 2]])
source = torch.tensor(
          [[[0, 1],
           [1, 2]],

           [[2, 3],
            [3, 4]]],
        dtype=torch.float,
        requires_grad=True,
)
result = sum_by_index(labels, source)

expected_output = torch.tensor([[0., 1.], [3., 5.], [3., 4.]], dtype=torch.float)

assert isinstance(result, torch.Tensor), 'result should be a torch.Tensor'

assert result.shape == expected_output.shape, f'result should have shape {expected_output.shape}, but got {result.shape}'

assert torch.allclose(result, expected_output), 'result does not match the expected output'

assert result.requires_grad, 'result should have requires_grad=True'

def sum_by_index(labels, source):
    unique_labels, _, labels_counts = tf.unique_with_counts(tf.reshape(labels, [-1]))
    labels_inds = tf.where(tf.equal(tf.expand_dims(labels, -1), tf.expand_dims(unique_labels, 0)))[..., -1]
    out = tf.zeros((tf.shape(unique_labels)[0], tf.shape(source)[-1]), dtype=source.dtype)

    # Add the embeddings from source to out based on labels_inds
    # out = ...

    return out

labels = tf.constant([[0, 1], [1, 2]])
source = tf.constant(
          [[[0, 1],
           [1, 2]],

           [[2, 3],
            [3, 4]]],
        dtype=tf.float32
)

# Apply sum_by_index to get the result
# result =

def sum_by_index(labels, source):
    unique_labels, _, labels_counts = tf.unique_with_counts(tf.reshape(labels, [-1]))
    labels_inds = tf.where(tf.equal(tf.expand_dims(labels, -1), tf.expand_dims(unique_labels, 0)))[..., -1]
    out = tf.zeros((tf.shape(unique_labels)[0], tf.shape(source)[-1]), dtype=source.dtype)
    out = tf.tensor_scatter_nd_add(out, tf.reshape(labels_inds, [-1, 1]), tf.reshape(source, [-1, tf.shape(source)[-1]]))
    return out

labels = tf.constant([[0, 1], [1, 2]])
source = tf.constant(
          [[[0, 1],
           [1, 2]],

           [[2, 3],
            [3, 4]]],
        dtype=tf.float32
)
result = sum_by_index(labels, source)

expected_output = tf.constant([[0., 1.], [3., 5.], [3., 4.]], dtype=tf.float32)

assert isinstance(result, tf.Tensor), 'result should be a tf.Tensor'

assert tf.reduce_all(tf.abs(result - expected_output) < 1e-6), 'result does not match the expected output'


original = torch.randn(size=(4, 3, 2))
row_exclude = 2

# Delete the row from x at index row_exclude
# x =

x = torch.cat((original[:row_exclude], original[row_exclude+1:]))

original_shape = (4, 3, 2)
expected_shape = (3, 3, 2)

assert isinstance(x, torch.Tensor), 'x should be a torch.Tensor'

assert x.shape == expected_shape, f'x should have shape {expected_shape}, but got {x.shape}'

assert torch.allclose(x[:row_exclude], original[:row_exclude]), 'Rows before row_exclude should be unchanged'

assert torch.allclose(x[row_exclude:], original[row_exclude+1:]), 'Rows after row_exclude should be shifted by 1'

original = tf.random.normal(shape=(4, 3, 2))
row_exclude = 2

# Delete the row from x at index row_exclude
# x =

x = tf.concat([original[:row_exclude], original[row_exclude+1:]], axis=0)

assert isinstance(x, tf.Tensor), 'x should be a tf.Tensor'

assert tf.reduce_all(x[:row_exclude] == original[:row_exclude]), 'Rows before row_exclude should be unchanged'

assert tf.reduce_all(x[row_exclude:] == original[row_exclude+1:]), 'Rows after row_exclude should be shifted by 1'


class TinyVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()

        # Define the convolutional block
        # self.conv_block1 = ...

        # Define the dropout layer
        # self.dropout = ...

        # Define the classifier block
        # self.classifier = ...

    def forward(self, x: torch.Tensor):
        # Forward pass through the convolutional block
        # x = ...

        # Apply dropout
        # x = ...

        # Forward pass through the classifier block
        # x = ...

        return x

class TinyVGG(nn.Module):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block1 = nn.Sequential(
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )

        self.dropout = nn.Dropout(0.4)
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(in_features=hidden_units*18*18, out_features=output_shape))

    def forward(self, x: torch.Tensor):
        x = self.conv_block1(x)
        x = self.dropout(x)
        x = self.classifier(x)
        return x

input_shape = 3
hidden_units = 64
output_shape = 10

model = TinyVGG(input_shape, hidden_units, output_shape)
sample_input = torch.randn(1, input_shape, 150, 150)

assert isinstance(model, nn.Module), 'TinyVGG should inherit from nn.Module'

assert isinstance(model.conv_block1, nn.Sequential), 'conv_block1 should be an instance of nn.Sequential'

assert isinstance(model.dropout, nn.Dropout), 'dropout should be an instance of nn.Dropout'

assert isinstance(model.classifier, nn.Sequential), 'classifier should be an instance of nn.Sequential'

output = model(sample_input)
assert output.shape == (1, output_shape), f'Output shape should be (1, {output_shape}), but got {output.shape}'

import tensorflow as tf
from tensorflow.keras import layers

class TinyVGG(tf.keras.Model):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()

        # Define the convolutional block
        # self.conv_block1 = ...

        # Define the dropout layer
        # self.dropout = ...

        # Define the classifier block
        # self.classifier = ...

    def call(self, x: tf.Tensor):
        # Forward pass through the convolutional block
        # x = ...

        # Apply dropout
        # x = ...

        # Forward pass through the classifier block
        # x = ...

        return x

class TinyVGG(tf.keras.Model):
    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        self.conv_block1 = tf.keras.Sequential([
            layers.Conv2D(filters=hidden_units, kernel_size=3, strides=1, padding='same', activation='relu', input_shape=(None, None, input_shape)),
            layers.MaxPooling2D(pool_size=2, strides=2),
            layers.Conv2D(filters=hidden_units, kernel_size=3, strides=1, padding='same', activation='relu'),
            layers.MaxPooling2D(pool_size=2, strides=2),
            layers.Conv2D(filters=hidden_units, kernel_size=3, strides=1, padding='same', activation='relu'),
            layers.MaxPooling2D(pool_size=2, strides=2)
        ])

        self.dropout = layers.Dropout(0.4)
        self.classifier = tf.keras.Sequential([layers.Flatten(), layers.Dense(units=output_shape)])

    def call(self, x: tf.Tensor):
        x = self.conv_block1(x)
        x = self.dropout(x)
        x = self.classifier(x)
        return x

input_shape = 3
hidden_units = 64
output_shape = 10

model = TinyVGG(input_shape, hidden_units, output_shape)
sample_input = tf.random.normal(shape=(1, 150, 150, input_shape))

assert isinstance(model, tf.keras.Model), 'TinyVGG should inherit from tf.keras.Model'

assert isinstance(model.conv_block1, tf.keras.Sequential), 'conv_block1 should be an instance of tf.keras.Sequential'

assert isinstance(model.dropout, layers.Dropout), 'dropout should be an instance of layers.Dropout'

assert isinstance(model.classifier, tf.keras.Sequential), 'classifier should be an instance of tf.keras.Sequential'


a = torch.Tensor(np.array([0.0917, -0.0006, 0.1825, -0.2484]))

# Create a one-hot tensor based on the index of the maximum value in a
# one_hot =

one_hot = torch.nn.functional.one_hot(a.argmax(), num_classes=a.numel())

expected_output = torch.tensor([0, 0, 1, 0])

assert isinstance(one_hot, torch.Tensor), 'one_hot should be a torch.Tensor'

assert one_hot.shape == a.shape, f'one_hot should have the same shape as a, but got {one_hot.shape}'

assert torch.sum(one_hot) == 1, 'one_hot should have only one non-zero element'

assert torch.all(one_hot == expected_output), f'one_hot does not match the expected output:\n{expected_output}'

a = tf.constant([0.0917, -0.0006, 0.1825, -0.2484], dtype=tf.float32)

# Create a one-hot tensor based on the index of the maximum value in a
# one_hot =

one_hot = tf.one_hot(tf.argmax(a), depth=tf.size(a))

import numpy as np
expected_output = tf.constant([0, 0, 1, 0], dtype=tf.int32)

assert isinstance(one_hot, tf.Tensor), 'one_hot should be a tf.Tensor'

assert tf.reduce_sum(one_hot) == 1, 'one_hot should have only one non-zero element'

assert np.all(one_hot.numpy() == expected_output.numpy()), f'one_hot does not match the expected output:\n{expected_output}'


import torch
import torchvision.transforms as T

# Compose the transformations
# trans_comp =

trans_comp = T.Compose([
  T.Resize([224, 224]),
  T.ToTensor()
])

from PIL import Image
import numpy as np
random_image_array = np.random.randint(0, 256, size=(256, 256, 3), dtype=np.uint8)
image = Image.fromarray(random_image_array)
transformed_image = trans_comp(image)

assert isinstance(transformed_image, torch.Tensor), 'transformed_image should be a torch.Tensor'

assert transformed_image.shape[1:] == (224, 224), f'transformed_image should have shape (C, 224, 224), but got {transformed_image.shape}'

assert transformed_image.min() >= 0 and transformed_image.max() <= 1, 'transformed_image should have values in the range [0, 1]'

# Define the resize function
# def resize(image):
#     return tf.image.resize(image, [224, 224])

# Define the normalize function
# def normalize(image):
#     return image / 255.0

# Apply the transformations to the dataset
# def transform(image, label):
#     image = resize(image)
#     image = normalize(image)
#     return image, label
#
# transformed_dataset = dataset.map(transform)

def resize(image):
    return tf.image.resize(image, [224, 224])

def normalize(image):
    return image / 255.0

def transform(image, label):
    image = resize(image)
    image = normalize(image)
    return image, label

import numpy as np

image = np.random.randint(0, 256, size=(256, 256, 3), dtype=np.uint8)
label = 0
image, label = transform(image, label)

assert isinstance(image, tf.Tensor), 'image should be a tf.Tensor'

assert image.shape == (224, 224, 3), f'image should have shape (224, 224, 3), but got {image.shape}'

assert tf.reduce_min(image) >= 0 and tf.reduce_max(image) <= 1, 'image should have values in the range [0, 1]'


a = torch.rand(8, 2)
b = torch.rand(8, 4)
c = torch.rand(8, 6)

# Concatenate the tensors along a new dimension
# d =

# Reshape the concatenated tensor into a 3D tensor
# e =

d = torch.cat((a, b, c), dim=1)
e = torch.reshape(d, (8, 3, -1))

assert isinstance(d, torch.Tensor), 'd should be a torch.Tensor'

assert d.shape == (8, 12), f'd should have shape (8, 12), but got {d.shape}'

assert isinstance(e, torch.Tensor), 'e should be a torch.Tensor'

assert e.shape == expected_shape, f'e should have shape {expected_shape}, but got {e.shape}'

a = tf.random.uniform((8, 2))
b = tf.random.uniform((8, 4))
c = tf.random.uniform((8, 6))

# Concatenate the tensors along a new dimension
# d =

# Reshape the concatenated tensor into a 3D tensor
# e =

d = tf.concat((a, b, c), axis=1)
e = tf.reshape(d, (8, 3, -1))

assert isinstance(d, tf.Tensor), 'd should be a tf.Tensor'

assert isinstance(e, tf.Tensor), 'e should be a tf.Tensor'


n = 5
a = torch.randn(1, 3, 8)  # Random (1, 3, 8) tensor

# Increase the first dimension to n and pad with zeros
# result =

result = torch.cat((a, torch.zeros(n - 1, a.shape[1], a.shape[2])))

assert isinstance(a, torch.Tensor), 'a should be a torch.Tensor'

assert a.shape == (1, 3, 8), f'a should have shape (1, 3, 8), but got {a.shape}'

assert isinstance(result, torch.Tensor), 'result should be a torch.Tensor'

assert result.shape == expected_shape, f'result should have shape {expected_shape}, but got {result.shape}'

assert torch.allclose(result[0], a), 'The first slice of result should be equal to a'

assert torch.allclose(result[1:], torch.zeros(n - 1, 3, 8)), 'The remaining slices of result should be zero'

n = 5
a = tf.random.normal((1, 3, 8))  # Random (1, 3, 8) tensor

# Increase the first dimension to n and pad with zeros
# result =

result = tf.concat((a, tf.zeros((n - 1, a.shape[1], a.shape[2]))), axis=0)

assert isinstance(a, tf.Tensor), 'a should be a tf.Tensor'

assert isinstance(result, tf.Tensor), 'result should be a tf.Tensor'

assert tf.reduce_all(tf.equal(result[0], a)), 'The first slice of result should be equal to a'

assert tf.reduce_all(tf.equal(result[1:], tf.zeros((n - 1, 3, 8)))), 'The remaining slices of result should be zero'


data = torch.arange(10)
starts = torch.tensor([0, 3, 4, 1])
ends = starts + 2

# Slice the data tensor using starts and ends
# result =

indices = torch.stack((starts, ends), axis=1)
result = torch.stack([data[slice(idx[0], idx[1])] for idx in indices])

assert result.shape == (4, 2), f'result should have shape (4, 2), but got {result.shape}'

assert torch.all(torch.eq(result, expected_result)), 'result does not match expected output'

data = tf.range(10)
starts = tf.constant([0, 3, 4, 1])
ends = starts + 2

# Slice the data tensor using starts and ends
# result =

indices = tf.stack([starts, ends], axis=1)
result = tf.stack([tf.slice(data, [start], [end - start]) for start, end in indices])

assert tf.reduce_all(tf.equal(result, expected_result)), 'result does not match expected output'


x = torch.tensor([1, 2, 3, 4])

# Subtract elements of every possible pair
# result =

assert result.shape == (4, 4), f'result should have shape (4, 4), but got {result.shape}'

x = tf.constant([1, 2, 3, 4])

# Subtract elements of every possible pair
# result =

result = x[:, tf.newaxis] - x[tf.newaxis, :]


a = torch.as_tensor([
    [[1, 3, 7, 6], [9, 0, 6, 2], [3, 0, 5, 8]],
    [[1, 0, 1, 0], [2, 1, 0, 3], [0, 0, 6, 1]]
])

sorted_dim = 1  # sort by rows, preserving each row
sorted_column = 2  # sort rows on value of 3rd column of each row

# Sort the tensor a row-wise based on the sorted_column
# sorted_a =

b = torch.argsort(a[:, :, sorted_column], dim=sorted_dim)
sorted_a = torch.stack([a[i, b[i], :] for i in range(a.shape[0])])

assert isinstance(sorted_a, torch.Tensor), 'sorted_a should be a torch.Tensor'

assert sorted_a.shape == a.shape, f'sorted_a should have the same shape as a, but got {sorted_a.shape} instead of {a.shape}'

assert torch.all(torch.eq(sorted_a, expected_result)), 'sorted_a does not match the expected result'

a = tf.constant([
    [[1, 3, 7, 6], [9, 0, 6, 2], [3, 0, 5, 8]],
    [[1, 0, 1, 0], [2, 1, 0, 3], [0, 0, 6, 1]]
])

sorted_dim = 1  # sort by rows, preserving each row
sorted_column = 2  # sort rows on value of 3rd column of each row

# Sort the tensor a row-wise based on the sorted_column
# sorted_a =

b = tf.argsort(a[:, :, sorted_column], axis=sorted_dim)
sorted_a = tf.stack([tf.gather(a[i], b[i], axis=0) for i in range(a.shape[0])])

assert isinstance(sorted_a, tf.Tensor), 'sorted_a should be a tf.Tensor'

assert tf.reduce_all(tf.equal(sorted_a, expected_result)), 'sorted_a does not match the expected result'


in_dict = {
    "train": [
        {
            "input": [[3, 1, 2], [3, 1, 2], [3, 1, 2]],
            "output": [[4, 5, 6], [4, 5, 6], [4, 5, 6]]
        },
        {
            "input": [[2, 3, 8], [2, 3, 8], [2, 3, 8]],
            "output": [[6, 4, 9], [6, 4, 9], [6, 4, 9]]
        }
    ]
}

# Convert the dictionary to a list of tensors
# train_examples = ...

# Apply padding to each input and output tensor
# for item in train_examples:
#     item[0] = ...
#     item[1] = ...

# Convert the tensors back to a dictionary
# out_dict = {'train': []}
# for item in train_examples:
#     out_dict['train'].append(...)

train_examples = []
for item in in_dict['train']:
    in_tensor = torch.Tensor(item['input'])
    out_tensor = torch.Tensor(item['output'])
    train_examples.append([in_tensor, out_tensor])

for item in train_examples:
    item[0] = F.pad(item[0], (1, 1, 1, 1))
    item[1] = F.pad(item[1], (1, 1, 1, 1))

out_dict = {'train': []}
for item in train_examples:
    out_dict['train'].append({
        'input': item[0].tolist(),
        'output': item[1].tolist()
    })

assert isinstance(train_examples, list), 'train_examples should be a list'

assert len(train_examples) == 2, f'Expected 2 training examples, got {len(train_examples)}'

assert all(isinstance(item, list) and len(item) == 2 for item in train_examples), 'Each item in train_examples should be a list of length 2'

assert all(isinstance(tensor, torch.Tensor) for item in train_examples for tensor in item), 'Each item in train_examples should contain torch.Tensor objects'

assert out_dict['train'][0]['input'] == expected_input[0], f"First input tensor does not match the expected padded tensor:\nExpected: {expected_input[0]}\nGot: {out_dict['train'][0]['input']}"

assert out_dict['train'][1]['input'] == expected_input[1], f"Second input tensor does not match the expected padded tensor:\nExpected: {expected_input[1]}\nGot: {out_dict['train'][1]['input']}"

assert out_dict['train'][0]['output'] == expected_output[0], f"First output tensor does not match the expected padded tensor:\nExpected: {expected_output[0]}\nGot: {out_dict['train'][0]['output']}"

assert out_dict['train'][1]['output'] == expected_output[1], f"Second output tensor does not match the expected padded tensor:\nExpected: {expected_output[1]}\nGot: {out_dict['train'][1]['output']}"

train_examples = []
for item in in_dict['train']:
    in_tensor = tf.convert_to_tensor(item['input'], dtype=tf.float32)
    out_tensor = tf.convert_to_tensor(item['output'], dtype=tf.float32)
    train_examples.append([in_tensor, out_tensor])

for item in train_examples:
    item[0] = tf.pad(item[0], [[1, 1], [1, 1]])
    item[1] = tf.pad(item[1], [[1, 1], [1, 1]])

out_dict = {'train': []}
for item in train_examples:
    out_dict['train'].append({
        'input': item[0].numpy().tolist(),
        'output': item[1].numpy().tolist()
    })

assert all(isinstance(tensor, tf.Tensor) for item in train_examples for tensor in item), 'Each item in train_examples should contain tf.Tensor objects'