| # Copyright (c) Microsoft Corporation. | |
| # Licensed under the MIT License. | |
| from typing import List | |
| import numpy as np | |
| def get_all_factors(n: int) -> List[int]: | |
| # Calculate the square root of n and round it up to the nearest integer | |
| n0 = int(np.ceil(np.sqrt(n))) | |
| # Find all divisors of n that are less than n0 | |
| val = np.where(n % np.arange(1, n0) == 0)[0] + 1 | |
| # If n is a perfect square, add the square root to the list of factors | |
| mid = np.array([], dtype=int) if n0 * n0 != n else [n0] | |
| # Combine the factors and their corresponding larger pair factors | |
| return [int(x) for x in np.concatenate([val, mid, n // val[::-1]])] | |
| def factorize(n: int) -> List[int]: | |
| i = 2 # Start with the smallest prime number | |
| result = [] | |
| # Iterate through numbers to find factors | |
| while n > 1: | |
| if n % i == 0: # If i is a factor of n | |
| n //= i # Divide n by i and keep the integer part | |
| result.append(i) | |
| else: | |
| i += 1 # Try the next number | |
| return result | |
| def coalesced_factor(subtensor: List[int], tensor: List[int]) -> int: | |
| # If the last dimension of the subtensor and tensor differ, or subtensor has only one dimension | |
| if subtensor[-1] != tensor[-1] or len(subtensor) == 1: | |
| return subtensor[-1] | |
| else: | |
| # Recursively calculate the coalesced factor for the remaining dimensions | |
| return subtensor[-1] * coalesced_factor(subtensor[:-1], tensor[:-1]) | |
| def coalesced_tensor_shape(subtensor: List[int], tensor: List[int], transaction_size: int) -> int: | |
| # Calculate the total number of elements in the subtensor | |
| bytes = int(np.prod(subtensor)) | |
| if bytes == 0: | |
| return 0 | |
| # Calculate the coalesced factor for the subtensor | |
| factor = int(coalesced_factor(subtensor, tensor)) | |
| # Compute the shape of the coalesced tensor | |
| return transaction_size * bytes / min(transaction_size, factor) | |