<jupyter_start><jupyter_text>Illustrate log-space computation in TFCode is derived fromhttps://github.com/EEA-sensors/sequential-parallelization-examples<jupyter_code>import math

import matplotlib.pyplot as plt
import numpy as np
import scipy as sc


from tensorflow import function
import tensorflow as tf
import tensorflow_probability as tfp

mm = tf.linalg.matmul
mv = tf.linalg.matvec
@tf.function
def log_mv(log_A, log_b, transpose_a=False):
    Amax = tf.reduce_max(log_A, axis=(-1, -2), keepdims=True)
    bmax = tf.reduce_max(log_b, axis=(-1), keepdims=True)
    return tf.squeeze(Amax, axis=-1) + bmax + tf.math.log(mv(tf.math.exp(log_A - Amax), tf.math.exp(log_b - bmax), transpose_a=transpose_a))

@tf.function
def semilog_mv(A, log_b, transpose_a=False):
    bmax = tf.reduce_max(log_b, axis=(-1), keepdims=True)
    return bmax + tf.math.log(mv(A, tf.math.exp(log_b - bmax), transpose_a=transpose_a))

@tf.function
def log_mm(log_A,log_B,transpose_a=False,transpose_b=False):
    Amax = tf.reduce_max(log_A, axis=(-1, -2), keepdims=True)
    Bmax = tf.reduce_max(log_B, axis=(-1, -2), keepdims=True)
    return Amax + Bmax + tf.math.log(mm(tf.math.exp(log_A - Amax), tf.math.exp(log_B - Bmax), transpose_a=transpose_a, transpose_b=transpose_b))

@tf.function
def log_normalize(log_p):
    pmax = tf.reduce_max(log_p, axis=(-1), keepdims=True)
    temp = tf.math.exp(log_p - pmax)
    return tf.math.log(temp / tf.reduce_sum(temp, axis=-1, keepdims=True))



print('Test utility functions:')
with tf.device('/CPU:0'):
    tf.random.set_seed(5)
    A = tf.random.uniform(shape=[4, 4])
    B = tf.random.uniform(shape=[4, 4])
    log_A = tf.math.log(A)
    log_B = tf.math.log(B)
    r1 = mm(A, B)
    r2 = tf.math.exp(log_mm(log_A, log_B))
    assert np.allclose(r1, r2)
    print(r1)

    b = tf.random.uniform(shape=[4])
    log_b = tf.math.log(b)
    r1 = mv(A, b)
    r2 = tf.math.exp(log_mv(log_A, log_b))
    r3 = tf.math.exp(semilog_mv(A, log_b))
    assert np.allclose(r1, r2)
    assert np.allclose(r1, r3)
    print(r1)


    r1 = b / tf.reduce_sum(b, keepdims=True)
    r2 = tf.math.exp(log_normalize(log_b))
    assert np.allclose(r1, r2)
    print(r1)<jupyter_output>Test utility functions:
tf.Tensor(
[[1.4257064  1.0391585  0.7607688  1.0484284 ]
 [0.9447347  0.76951545 0.44059256 0.6164023 ]
 [0.35811597 0.34018368 0.16159499 0.33705997]
 [0.75753486 0.5810969  0.41004145 0.6656992 ]], shape=(4, 4), dtype=float32)
tf.Tensor([0.5768026  0.33639905 0.15708235 0.353592  ], shape=(4,), dtype=float32)
tf.Tensor([0.19298144 0.3319118  0.0531239  0.42198285], shape=(4,), dtype=float32)
{'fraction_non_alphanumeric': 0.1557596632223498, 'fraction_numerical': 0.09529276693455797, 'mean_word_length': 4.075728155339806, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '522709', 'license': 'MIT', 'n_tokens_mistral': 1261, 'n_tokens_neox': 1141, 'n_words': 202, 'path': 'notebooks/logspace_tf.ipynb', 'repo_name': 'arpitvaghela/probml-notebooks', 'size': '2613'}