#!/usr/bin/env python3

import torch

nMatrix = 50
nRow    = 40
nCol    = 30
nS      = 20

AA = torch.randn(nMatrix,nRow,nCol)             # combine many matrices
BB = torch.randn(nMatrix,nCol,nS  )             # to a single tensor

YY = torch.zeros(nMatrix,nRow,nS)   
if (1==1): 
  YY = torch.matmul(AA,BB)                      # tensor operation
else:
  for ii in range(nMatrix):                     # looping over all matrices
    for nn in range(nRow):          
      for mm in range(nS):          
        for ll in range(nCol):                  # explicit multiplication
           YY[ii][nn][mm] += AA[ii][nn][ll]*BB[ii][ll][mm]
#
print()
print("AA", AA.shape)
print("BB", BB.shape)
print("YY", YY.shape)
print()
print(f'we did multiply {nMatrix:d} matrices')
print(f'of type ({nRow:d}x{nCol:d}) and ({nCol:d}x{nS:d})')

#!/usr/bin/env python3

"""
torch.vmap() vs. manual looping example

Some alorithms need per-sample linear transforms:
For a batch of inputs x_i, each sample has its 
own weight matrix W_i and bias b_i:

  y_i = W_i @ x_i + b_i    for i = 0 ... batch_size-1

A plain `torch.matmul()` can not do this in one 
shot, because normal batched matmul broadcasts 
shared weights, not different weights per sample.

This is a situation for torch.vmap(): 
take a function written for a single 
example, and vectorize it without for-loops.
"""

import torch
torch.manual_seed(0)   # just for fun

batch_size = 4
in_dim     = 3
out_dim    = 2

# per sample weight matrices W (and bias b)
W = torch.randn(batch_size, out_dim, in_dim) 
b = torch.randn(batch_size, out_dim)  
# input
x = torch.randn(batch_size, in_dim)      

# ---------------------------
# linear transform for a for a single example
# ---------------------------

# no batch dimension involved; with hints
def linear_single(w_i: torch.Tensor, 
                  b_i: torch.Tensor, 
                  x_i: torch.Tensor) -> torch.Tensor:
  return w_i @ x_i + b_i

# ---------------------------
# witout vmap(): manual looping over batch dimension
# ---------------------------

outputs_loop = []
for i in range(batch_size):
    y_i = linear_single(W[i], b[i], x[i])
    outputs_loop.append(y_i)
# stack: list --> tensor 
# shape: (batch_size, out_dim)
y_loop = torch.stack(outputs_loop, dim=0)

# ---------------------------
# with vmap()
# ---------------------------

# using vmap(), create batched version of 'linear_single()'
#
# in_dims=(0, 0, 0):
# batch dimension is 0 for all three arguments:
# w_i, b_i, x_i 

batched_linear = torch.vmap(linear_single, in_dims=(0, 0, 0))

# shape: (batch_size, out_dim)
y_vmap = batched_linear(W, b, x)      

# ---------------------------
# identical results?
# ---------------------------

print("Loop result:\n", y_loop)
print("\nvmap result:\n", y_vmap)
print("\nMatch:", torch.allclose(y_loop, y_vmap))

#!/usr/bin/env python3

# conditional tensor operations

import torch

xIn  = torch.randn(6)
xOut = torch.where(xIn > 0, 1.0, 0.0)     # conditional mapping
print("xIn \n", xIn)
print("torch.where(xIn > 0, 1.0, 0.0)\n", xOut)
print()

yIn    = torch.arange(10)
yOnes  = torch.ones(10).int()             # float to int
yOut = torch.where(yIn%2==0, yIn, yOnes)  # conditional substitution

print("yIn   \n", yIn)
print("yOnes \n", yOnes)
print("torch.where(yIn%2==0, yIn, yOnes)\n", yOut)
print()

# element-wise stochastic operations
# need random masks
zIn      = torch.arange(6)
randR    = torch.rand(6)               # in [0,1]
randMask = torch.where(randR<0.5,1,0)  # either 0 or 1
zOut = torch.where(zIn%2==0, zIn, randMask)

print("zIn      \n", zIn)
print("randR    \n", randR)
print("randMask \n", randMask)
print("torch.where(zIn%2==0, zIn, randMask)\n", zOut)

#!/usr/bin/env python3

import threading   # nomen est omen
import time        # sleep, etc.

#
# function to run asynchronous
#
def do_work(tNumber=0):
    print("starting  thread # %s", tNumber)
    rr = 1.0
    if (tNumber==0):
      time.sleep(10)             # sleeping 10 seconds
    else:
      while (1==1):              # doing heavy stuff
        rr = 1.0/(1.0+rr)
    print("finishing thread # %s", tNumber)

#
# main
#
allThreads = list()
for i in range(5):
  x = threading.Thread(target=do_work, args=(i,))
  allThreads.append(x)
  x.start()                      # starting thread 
#
for i in range(len(allThreads)):
  allThreads[i].join()           # waiting for threads to finish
#
print("\n# all done folks")

#!/usr/bin/env python3

import torch
import os                          # operating system

nMatrix = 500
nRow    = 400
nCol    = 300
nS      = 200

torch.set_num_threads(1)           # for cpu-based hardware
nCPU = os.cpu_count()              # number of available cpu
if (nCPU>2): 
  torch.set_num_threads(nCPU-2)    # leave two for other uses
#
print("\n# number of CPU, threads: ",nCPU,torch.get_num_threads()) 
#
AA = torch.randn(nMatrix,nRow,nCol)            
BB = torch.randn(nMatrix,nCol,nS  )        
#
for ii in range(1000):             # do heavy stuff
  if (ii%50==0):
    print(ii)
  YY = torch.matmul(AA,BB)    

#!/usr/bin/env python3

import os                         # underlying operation system

aa = 10
print()
print(f"main PID:", os.getpid())  # pid of main
print()

newPid = os.fork()                # forking 
myPid = os.getpid()               # pid after forking

aa = aa + 2
print(f"   myPID: {myPid:6d},  newPid: {newPid:6d},                    aa: {aa:6d}")

newNewPid = os.fork()
myPid = os.getpid()

aa = aa + 10
print(f"   myPID: {myPid:6d},  newPid: {newPid:6d}, newNewPid: {newNewPid:6d}, aa: {aa:6d}")

#!/usr/bin/env python3

# simple example of a functional network
# size defined by arguments

# leaky_relu(input, negative_slope=0.01) 

import torch
import torch.nn.functional as F

def functionalNetwork(x, w1, w2):
  h = F.linear(x, w1)
# h = F.relu(h)
  h = F.leaky_relu(h)
  y = F.linear(h, w2)
  return F.softmax(y, dim=1)

# usage
x  = torch.randn(64, 1000)   # 64 samples, 1000 features
w1 = torch.randn(500, 1000)  # first layer weights
w2 = torch.randn(10, 500)    # second layer weights

output = functionalNetwork(x, w1, w2)