view release on metacpan or  search on metacpan

t/test_kvstore.t  view on Meta::CPAN

use strict;
use warnings;
use Test::More tests => 38;
use AI::MXNet qw(mx);

my $shape = [4, 4];
my $keys  = [5,7,9];

sub init_kv
{
    # init kv
    my $kv = mx->kv->create();
    # single
    $kv->init(3, mx->nd->zeros($shape));
    # list
    $kv->init($keys, [map { mx->nd->zeros($shape) } 0..@$keys-1]);
    return $kv;
}

sub check_diff_to_scalar
{
    # assert A == x
    my ($A, $x) = @_;
    ok(($A - $x)->aspdl->abs->sum == 0);
}

sub test_single_kv_pair
{
    # single key-value pair push & pull
    my $kv = init_kv();
    $kv->push(3, mx->nd->ones($shape));
    my $val = mx->nd->empty($shape);
    $kv->pull(3, out => $val);
    check_diff_to_scalar($val, 1);
}

sub test_init
{
    my $kv = mx->kv->create();
    $kv->init(3, mx->nd->ones($shape)*4);
    my $a = mx->nd->zeros($shape);
    $kv->pull(3, out=>$a);
    check_diff_to_scalar($a, 4);
}

sub test_list_kv_pair
{
    # list key-value pair push & pull
    my $kv = init_kv();
    $kv->push($keys, [map {mx->nd->ones($shape)*4} 0..@$keys-1]);
    my $val = [map { mx->nd->empty($shape) } 0..@$keys-1];
    $kv->pull($keys, out => $val);
    for my $v (@$val)
    {
        check_diff_to_scalar($v, 4);
    }
}

sub test_aggregator
{
    # aggregate value on muliple devices

    my $kv = init_kv();