import tensorflow as tf 
def ARelu(x, alpha=0.90, beta=2.0):
    alpha = tf.clip_by_value(alpha, clip_value_min=0.01, clip_value_max=0.99)
    beta  = 1 + tf.math.sigmoid(beta)
    return tf.nn.relu(x) * beta - tf.nn.relu(-x) * alpha

import tensorflow as tf
from tensorflow import keras
import numpy as np

(xtrain, ytrain), (xtest, ytest) = keras.datasets.mnist.load_data()

xtrain = np.float32(xtrain/255)
xtest = np.float32(xtest/255)

ytrain = np.int32(ytrain)
ytest = np.int32(ytest)

def pre_process(inputs, targets):
    inputs = tf.expand_dims(inputs, -1)
    targets = tf.one_hot(targets, depth=10)
    return tf.divide(inputs, 255), targets

train_data = tf.data.Dataset.from_tensor_slices((xtrain, ytrain)).\
    take(10_000).shuffle(10_000).batch(8).map(pre_process)
test_data = tf.data.Dataset.from_tensor_slices((xtest, ytest)).\
    take(1_000).shuffle(1_000).batch(8).map(pre_process)

model = tf.keras.Sequential([
            tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), strides=(1, 1), 
                   input_shape=(28, 28, 1), activation='relu'),
            tf.keras.layers.MaxPool2D(pool_size=(2, 2)),
            tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), strides=(1, 1), 
                   activation=ARelu),
            tf.keras.layers.MaxPool2D(pool_size=(2, 2)),
            tf.keras.layers.Flatten(),
            tf.keras.layers.Dense(64, activation=ARelu), 
            tf.keras.layers.Dropout(5e-1),
            tf.keras.layers.Dense(10, activation='softmax')])

model.compile(loss='categorical_crossentropy', optimizer='adam')

history = model.fit(train_data, validation_data=test_data, epochs=10)